Files
2025-12-02 16:37:00 -08:00

659 lines
21 KiB
Python

"""
Channel Utilities
Pack and unpack multiple maps into RGB channels
"""
import torch
class ChannelPacker:
"""
Pack 3 grayscale maps into RGB channels
Common: Pack Roughness (R), Metallic (G), AO (B)
"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"red_channel": ("IMAGE",),
},
"optional": {
"green_channel": ("IMAGE",),
"blue_channel": ("IMAGE",),
"preset": (["custom", "orm_unity", "orm_unreal", "rma"], {
"default": "custom",
"tooltip": "Preset channel layouts: ORM (AO/Roughness/Metallic), RMA (Roughness/Metallic/AO)"
}),
}
}
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("packed",)
FUNCTION = "pack"
CATEGORY = "Texture Alchemist/Channel"
def pack(self, red_channel, green_channel=None, blue_channel=None, preset="custom"):
"""Pack grayscale images into RGB channels"""
print("\n" + "="*60)
print("Channel Packer")
print("="*60)
print(f"Preset: {preset}")
# Get dimensions from red channel
batch, height, width, channels = red_channel.shape
device = red_channel.device
dtype = red_channel.dtype
# Convert to grayscale if needed
r = self._to_grayscale(red_channel)
g = self._to_grayscale(green_channel) if green_channel is not None else torch.ones_like(r)
b = self._to_grayscale(blue_channel) if blue_channel is not None else torch.ones_like(r)
# Resize to match if sizes differ
if green_channel is not None:
g = self._resize_to_match(g, r)
if blue_channel is not None:
b = self._resize_to_match(b, r)
# Pack into RGB
packed = torch.cat([r, g, b], dim=-1)
print(f"✓ Channels packed")
print(f" R: {r.shape}")
print(f" G: {g.shape}")
print(f" B: {b.shape}")
print(f" Output: {packed.shape}")
print("="*60 + "\n")
return (packed,)
def _to_grayscale(self, image):
"""Convert to single channel grayscale"""
if image.shape[-1] == 1:
return image
elif image.shape[-1] == 3:
# Luminance
weights = torch.tensor([0.2989, 0.5870, 0.1140],
device=image.device, dtype=image.dtype)
return torch.sum(image * weights, dim=-1, keepdim=True)
else:
# Just take first channel
return image[:, :, :, 0:1]
def _resize_to_match(self, source, target):
"""Resize source to match target dimensions"""
if source.shape[1:3] == target.shape[1:3]:
return source
target_h, target_w = target.shape[1:3]
source_bchw = source.permute(0, 3, 1, 2)
resized = torch.nn.functional.interpolate(
source_bchw,
size=(target_h, target_w),
mode='bilinear',
align_corners=False
)
return resized.permute(0, 2, 3, 1)
class ChannelUnpacker:
"""
Unpack RGB channels into separate grayscale maps
"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"packed": ("IMAGE",),
"output_channels": (["all", "r_only", "g_only", "b_only", "rg", "rb", "gb"], {
"default": "all",
"tooltip": "Which channels to extract"
}),
}
}
RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE")
RETURN_NAMES = ("red", "green", "blue")
FUNCTION = "unpack"
CATEGORY = "Texture Alchemist/Channel"
def unpack(self, packed, output_channels="all"):
"""Unpack RGB channels into separate images"""
print("\n" + "="*60)
print("Channel Unpacker")
print("="*60)
print(f"Input shape: {packed.shape}")
print(f"Output channels: {output_channels}")
batch, height, width, channels = packed.shape
device = packed.device
dtype = packed.dtype
# Extract channels
if channels >= 3:
r = packed[:, :, :, 0:1]
g = packed[:, :, :, 1:2]
b = packed[:, :, :, 2:3]
elif channels == 2:
r = packed[:, :, :, 0:1]
g = packed[:, :, :, 1:2]
b = torch.zeros((batch, height, width, 1), device=device, dtype=dtype)
else:
r = packed
g = torch.zeros((batch, height, width, 1), device=device, dtype=dtype)
b = torch.zeros((batch, height, width, 1), device=device, dtype=dtype)
# Convert to RGB for display
r_rgb = r.repeat(1, 1, 1, 3)
g_rgb = g.repeat(1, 1, 1, 3)
b_rgb = b.repeat(1, 1, 1, 3)
print(f"✓ Channels unpacked")
print(f" Red: {r_rgb.shape}")
print(f" Green: {g_rgb.shape}")
print(f" Blue: {b_rgb.shape}")
print("="*60 + "\n")
return (r_rgb, g_rgb, b_rgb)
class GrayscaleToColor:
"""
Convert grayscale image to RGB by repeating the channel
"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"grayscale": ("IMAGE",),
}
}
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("color",)
FUNCTION = "convert"
CATEGORY = "Texture Alchemist/Channel"
def convert(self, grayscale):
"""Convert grayscale to RGB"""
print("\n" + "="*60)
print("Grayscale to Color")
print("="*60)
print(f"Input shape: {grayscale.shape}")
batch, height, width, channels = grayscale.shape
# If already RGB or more, return as-is
if channels >= 3:
print("✓ Already color (RGB+), returning as-is")
print("="*60 + "\n")
return (grayscale,)
# Convert to grayscale first if needed
if channels == 1:
gray = grayscale
else:
# Average multiple channels
gray = torch.mean(grayscale, dim=-1, keepdim=True)
# Repeat to create RGB
color = gray.repeat(1, 1, 1, 3)
print(f"✓ Converted to color")
print(f" Output shape: {color.shape}")
print("="*60 + "\n")
return (color,)
class ColorToGrayscale:
"""
Convert color (RGB) image to grayscale using luminance
"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"color": ("IMAGE",),
"method": (["luminance", "average", "lightness", "red_only", "green_only", "blue_only"], {
"default": "luminance",
"tooltip": "Conversion method: luminance (perceptual), average, lightness (min+max/2), or single channel"
}),
}
}
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("grayscale",)
FUNCTION = "convert"
CATEGORY = "Texture Alchemist/Channel"
def convert(self, color, method="luminance"):
"""Convert color to grayscale"""
print("\n" + "="*60)
print("Color to Grayscale")
print("="*60)
print(f"Input shape: {color.shape}")
print(f"Method: {method}")
batch, height, width, channels = color.shape
device = color.device
dtype = color.dtype
# If already grayscale, return as-is (as RGB for display)
if channels == 1:
print("✓ Already grayscale, returning as RGB")
gray = color.repeat(1, 1, 1, 3)
print("="*60 + "\n")
return (gray,)
# Extract RGB channels (handle RGBA by ignoring alpha)
r = color[:, :, :, 0:1]
g = color[:, :, :, 1:2] if channels >= 2 else r
b = color[:, :, :, 2:3] if channels >= 3 else r
# Apply conversion method
if method == "luminance":
# Perceptual luminance weights (Rec. 709)
weights = torch.tensor([0.2126, 0.7152, 0.0722], device=device, dtype=dtype)
gray_1ch = r * weights[0] + g * weights[1] + b * weights[2]
elif method == "average":
# Simple average
gray_1ch = (r + g + b) / 3.0
elif method == "lightness":
# (max + min) / 2
rgb_stack = torch.cat([r, g, b], dim=-1)
gray_1ch = (torch.max(rgb_stack, dim=-1, keepdim=True)[0] +
torch.min(rgb_stack, dim=-1, keepdim=True)[0]) / 2.0
elif method == "red_only":
gray_1ch = r
elif method == "green_only":
gray_1ch = g
elif method == "blue_only":
gray_1ch = b
else:
gray_1ch = r * 0.2126 + g * 0.7152 + b * 0.0722
# Convert to RGB for display
grayscale = gray_1ch.repeat(1, 1, 1, 3)
print(f"✓ Converted to grayscale")
print(f" Output shape: {grayscale.shape}")
print("="*60 + "\n")
return (grayscale,)
class ChannelPackerORMA:
"""
Pack 4 maps into RGBA: Occlusion (R), Roughness (G), Metallic (B), Alpha (A)
Advanced ORM with alpha channel support
"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"occlusion": ("IMAGE",),
"roughness": ("IMAGE",),
"metallic": ("IMAGE",),
},
"optional": {
"alpha": ("IMAGE",),
}
}
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("orma_packed",)
FUNCTION = "pack"
CATEGORY = "Texture Alchemist/Channel"
def pack(self, occlusion, roughness, metallic, alpha=None):
"""Pack ORMA channels"""
print("\n" + "="*60)
print("Channel Packer (ORMA)")
print("="*60)
# Get dimensions from occlusion
batch, height, width, channels = occlusion.shape
device = occlusion.device
dtype = occlusion.dtype
# Convert to grayscale if needed
o = self._to_grayscale(occlusion)
r = self._to_grayscale(roughness)
m = self._to_grayscale(metallic)
# Resize to match if sizes differ
r = self._resize_to_match(r, o)
m = self._resize_to_match(m, o)
# Handle alpha channel
if alpha is not None:
a = self._to_grayscale(alpha)
a = self._resize_to_match(a, o)
else:
# Default to fully opaque
a = torch.ones_like(o)
# Pack into RGBA
packed = torch.cat([o, r, m, a], dim=-1)
print(f"✓ ORMA packed")
print(f" Occlusion (R): {o.shape}")
print(f" Roughness (G): {r.shape}")
print(f" Metallic (B): {m.shape}")
print(f" Alpha (A): {a.shape}")
print(f" Output (RGBA): {packed.shape}")
print("="*60 + "\n")
return (packed,)
def _to_grayscale(self, image):
"""Convert to single channel grayscale"""
if image.shape[-1] == 1:
return image
elif image.shape[-1] >= 3:
# Luminance
weights = torch.tensor([0.2989, 0.5870, 0.1140],
device=image.device, dtype=image.dtype)
rgb = image[:, :, :, 0:3]
return torch.sum(rgb * weights, dim=-1, keepdim=True)
else:
# Just take first channel
return image[:, :, :, 0:1]
def _resize_to_match(self, source, target):
"""Resize source to match target dimensions"""
if source.shape[1:3] == target.shape[1:3]:
return source
target_h, target_w = target.shape[1:3]
source_bchw = source.permute(0, 3, 1, 2)
resized = torch.nn.functional.interpolate(
source_bchw,
size=(target_h, target_w),
mode='bilinear',
align_corners=False
)
return resized.permute(0, 2, 3, 1)
class ChannelPackerRMA:
"""
Pack RMA: Roughness (R), Metallic (G), AO (B)
Alternative to ORM format
"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"roughness": ("IMAGE",),
"metallic": ("IMAGE",),
"ao": ("IMAGE",),
}
}
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("rma_packed",)
FUNCTION = "pack"
CATEGORY = "Texture Alchemist/Channel"
def pack(self, roughness, metallic, ao):
"""Pack RMA channels"""
print("\n" + "="*60)
print("Channel Packer (RMA)")
print("="*60)
# Get dimensions from roughness
batch, height, width, channels = roughness.shape
device = roughness.device
dtype = roughness.dtype
# Convert to grayscale if needed
r = self._to_grayscale(roughness)
m = self._to_grayscale(metallic)
a = self._to_grayscale(ao)
# Resize to match if sizes differ
m = self._resize_to_match(m, r)
a = self._resize_to_match(a, r)
# Pack into RGB
packed = torch.cat([r, m, a], dim=-1)
print(f"✓ RMA packed")
print(f" Roughness (R): {r.shape}")
print(f" Metallic (G): {m.shape}")
print(f" AO (B): {a.shape}")
print(f" Output (RGB): {packed.shape}")
print("="*60 + "\n")
return (packed,)
def _to_grayscale(self, image):
"""Convert to single channel grayscale"""
if image.shape[-1] == 1:
return image
elif image.shape[-1] >= 3:
# Luminance
weights = torch.tensor([0.2989, 0.5870, 0.1140],
device=image.device, dtype=image.dtype)
rgb = image[:, :, :, 0:3]
return torch.sum(rgb * weights, dim=-1, keepdim=True)
else:
# Just take first channel
return image[:, :, :, 0:1]
def _resize_to_match(self, source, target):
"""Resize source to match target dimensions"""
if source.shape[1:3] == target.shape[1:3]:
return source
target_h, target_w = target.shape[1:3]
source_bchw = source.permute(0, 3, 1, 2)
resized = torch.nn.functional.interpolate(
source_bchw,
size=(target_h, target_w),
mode='bilinear',
align_corners=False
)
return resized.permute(0, 2, 3, 1)
class ChannelPackerRMAA:
"""
Pack RMAA: Roughness (R), Metallic (G), AO (B), Alpha (A)
RMA format with alpha channel support
"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"roughness": ("IMAGE",),
"metallic": ("IMAGE",),
"ao": ("IMAGE",),
},
"optional": {
"alpha": ("IMAGE",),
}
}
RETURN_TYPES = ("IMAGE",)
RETURN_NAMES = ("rmaa_packed",)
FUNCTION = "pack"
CATEGORY = "Texture Alchemist/Channel"
def pack(self, roughness, metallic, ao, alpha=None):
"""Pack RMAA channels"""
print("\n" + "="*60)
print("Channel Packer (RMAA)")
print("="*60)
# Get dimensions from roughness
batch, height, width, channels = roughness.shape
device = roughness.device
dtype = roughness.dtype
# Convert to grayscale if needed
r = self._to_grayscale(roughness)
m = self._to_grayscale(metallic)
a = self._to_grayscale(ao)
# Resize to match if sizes differ
m = self._resize_to_match(m, r)
a = self._resize_to_match(a, r)
# Handle alpha channel
if alpha is not None:
alpha_ch = self._to_grayscale(alpha)
alpha_ch = self._resize_to_match(alpha_ch, r)
else:
# Default to fully opaque
alpha_ch = torch.ones_like(r)
# Pack into RGBA
packed = torch.cat([r, m, a, alpha_ch], dim=-1)
print(f"✓ RMAA packed")
print(f" Roughness (R): {r.shape}")
print(f" Metallic (G): {m.shape}")
print(f" AO (B): {a.shape}")
print(f" Alpha (A): {alpha_ch.shape}")
print(f" Output (RGBA): {packed.shape}")
print("="*60 + "\n")
return (packed,)
def _to_grayscale(self, image):
"""Convert to single channel grayscale"""
if image.shape[-1] == 1:
return image
elif image.shape[-1] >= 3:
# Luminance
weights = torch.tensor([0.2989, 0.5870, 0.1140],
device=image.device, dtype=image.dtype)
rgb = image[:, :, :, 0:3]
return torch.sum(rgb * weights, dim=-1, keepdim=True)
else:
# Just take first channel
return image[:, :, :, 0:1]
def _resize_to_match(self, source, target):
"""Resize source to match target dimensions"""
if source.shape[1:3] == target.shape[1:3]:
return source
target_h, target_w = target.shape[1:3]
source_bchw = source.permute(0, 3, 1, 2)
resized = torch.nn.functional.interpolate(
source_bchw,
size=(target_h, target_w),
mode='bilinear',
align_corners=False
)
return resized.permute(0, 2, 3, 1)
class ChannelUnpackerRMA:
"""
Unpack RMA: Extract Roughness, Metallic, AO from packed RGB
"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"rma_packed": ("IMAGE",),
}
}
RETURN_TYPES = ("IMAGE", "IMAGE", "IMAGE")
RETURN_NAMES = ("roughness", "metallic", "ao")
FUNCTION = "unpack"
CATEGORY = "Texture Alchemist/Channel"
def unpack(self, rma_packed):
"""Unpack RMA channels"""
print("\n" + "="*60)
print("Channel Unpacker (RMA)")
print("="*60)
print(f"Input shape: {rma_packed.shape}")
batch, height, width, channels = rma_packed.shape
device = rma_packed.device
dtype = rma_packed.dtype
# Extract channels
if channels >= 3:
roughness = rma_packed[:, :, :, 0:1]
metallic = rma_packed[:, :, :, 1:2]
ao = rma_packed[:, :, :, 2:3]
elif channels == 2:
roughness = rma_packed[:, :, :, 0:1]
metallic = rma_packed[:, :, :, 1:2]
ao = torch.zeros((batch, height, width, 1), device=device, dtype=dtype)
else:
roughness = rma_packed
metallic = torch.zeros((batch, height, width, 1), device=device, dtype=dtype)
ao = torch.zeros((batch, height, width, 1), device=device, dtype=dtype)
# Convert to RGB for display
roughness_rgb = roughness.repeat(1, 1, 1, 3)
metallic_rgb = metallic.repeat(1, 1, 1, 3)
ao_rgb = ao.repeat(1, 1, 1, 3)
print(f"✓ RMA unpacked")
print(f" Roughness: {roughness_rgb.shape}")
print(f" Metallic: {metallic_rgb.shape}")
print(f" AO: {ao_rgb.shape}")
print("="*60 + "\n")
return (roughness_rgb, metallic_rgb, ao_rgb)
# Node registration
NODE_CLASS_MAPPINGS = {
"ChannelPacker": ChannelPacker,
"ChannelUnpacker": ChannelUnpacker,
"GrayscaleToColor": GrayscaleToColor,
"ColorToGrayscale": ColorToGrayscale,
"ChannelPackerORMA": ChannelPackerORMA,
"ChannelPackerRMA": ChannelPackerRMA,
"ChannelPackerRMAA": ChannelPackerRMAA,
"ChannelUnpackerRMA": ChannelUnpackerRMA,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"ChannelPacker": "Channel Packer (RGB)",
"ChannelUnpacker": "Channel Unpacker (RGB)",
"GrayscaleToColor": "Grayscale to Color",
"ColorToGrayscale": "Color to Grayscale",
"ChannelPackerORMA": "Channel Packer (ORMA)",
"ChannelPackerRMA": "Channel Packer (RMA)",
"ChannelPackerRMAA": "Channel Packer (RMAA)",
"ChannelUnpackerRMA": "Channel Unpacker (RMA)",
}