93 lines
3.6 KiB
Python
93 lines
3.6 KiB
Python
from PIL import Image
|
|
import numpy as np
|
|
import torch
|
|
|
|
def tensor_to_pil(tensor):
|
|
"""Convert ComfyUI tensor to PIL Image"""
|
|
# ComfyUI tensors are typically in format [batch, height, width, channels]
|
|
if len(tensor.shape) == 4:
|
|
tensor = tensor.squeeze(0) # Remove batch dimension
|
|
|
|
# Convert from tensor to numpy array and scale to 0-255
|
|
np_image = (tensor.cpu().numpy() * 255).astype(np.uint8)
|
|
|
|
# Handle different channel counts
|
|
if np_image.shape[2] == 4: # RGBA
|
|
return Image.fromarray(np_image, 'RGBA')
|
|
elif np_image.shape[2] == 3: # RGB
|
|
return Image.fromarray(np_image, 'RGB')
|
|
else: # Grayscale or other
|
|
return Image.fromarray(np_image[:,:,0], 'L').convert('RGBA')
|
|
|
|
def pil_to_tensor(pil_image):
|
|
"""Convert PIL Image to ComfyUI tensor"""
|
|
# Convert to RGB for output (ComfyUI typically expects RGB)
|
|
if pil_image.mode == 'RGBA':
|
|
# Create a white background and composite the RGBA image onto it
|
|
background = Image.new('RGB', pil_image.size, (255, 255, 255))
|
|
background.paste(pil_image, mask=pil_image.split()[3]) # Use alpha channel as mask
|
|
pil_image = background
|
|
elif pil_image.mode != 'RGB':
|
|
pil_image = pil_image.convert('RGB')
|
|
|
|
np_image = np.array(pil_image).astype(np.float32) / 255.0
|
|
tensor = torch.from_numpy(np_image)
|
|
return tensor.unsqueeze(0) # Add batch dimension
|
|
|
|
class MultiAlphaComposite:
|
|
CATEGORY = "Image/Composite"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"num_images": ("INT", {"default": 2, "min": 2, "max": 8, "step": 1}),
|
|
"image1": ("IMAGE",),
|
|
"image2": ("IMAGE",),
|
|
},
|
|
"optional": {
|
|
"image3": ("IMAGE",),
|
|
"image4": ("IMAGE",),
|
|
"image5": ("IMAGE",),
|
|
"image6": ("IMAGE",),
|
|
"image7": ("IMAGE",),
|
|
"image8": ("IMAGE",),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "composite"
|
|
|
|
def composite(self, num_images, image1, image2, image3=None, image4=None, image5=None, image6=None, image7=None, image8=None):
|
|
# Collect all provided images up to num_images
|
|
images = [image1, image2, image3, image4, image5, image6, image7, image8]
|
|
valid_images = [img for img in images[:num_images] if img is not None]
|
|
|
|
if len(valid_images) < 2:
|
|
raise ValueError("At least 2 images are required for compositing")
|
|
|
|
print(f"Debug: Processing {len(valid_images)} images for layer compositing")
|
|
|
|
# Convert ComfyUI tensors to PIL Images and ensure RGBA
|
|
layers = []
|
|
for i, img_tensor in enumerate(valid_images):
|
|
pil_img = tensor_to_pil(img_tensor)
|
|
# Convert to RGBA for alpha compositing (like Photoshop layers)
|
|
rgba_layer = pil_img.convert('RGBA')
|
|
layers.append(rgba_layer)
|
|
print(f"Debug: Layer {i+1} - Size: {rgba_layer.size}, Mode: {rgba_layer.mode}")
|
|
|
|
# Start with the bottom layer (image1)
|
|
result = layers[0]
|
|
print(f"Debug: Starting with base layer (image1)")
|
|
|
|
# Stack each subsequent layer on top (like Photoshop)
|
|
for i in range(1, len(layers)):
|
|
print(f"Debug: Compositing layer {i+1} on top")
|
|
result = Image.alpha_composite(result, layers[i])
|
|
|
|
print(f"Debug: Final composite size: {result.size}, mode: {result.mode}")
|
|
|
|
# Convert back to ComfyUI tensor format
|
|
result_tensor = pil_to_tensor(result)
|
|
return (result_tensor,) |