125 lines
3.7 KiB
Python
125 lines
3.7 KiB
Python
"""
|
|
Example usage script for ComfyUI Universal Image Loader
|
|
|
|
This script demonstrates how to use the ImageLoader node programmatically.
|
|
"""
|
|
|
|
import base64
|
|
import io
|
|
from PIL import Image
|
|
import torch
|
|
|
|
# Import the node
|
|
from imgloader_node import ImageLoader
|
|
|
|
|
|
def create_test_image():
|
|
"""Create a simple test image for demonstration."""
|
|
# Create a simple gradient image
|
|
img = Image.new('RGB', (200, 200))
|
|
pixels = img.load()
|
|
|
|
for i in range(200):
|
|
for j in range(200):
|
|
pixels[j, i] = (i, j, (i + j) % 255)
|
|
|
|
return img
|
|
|
|
|
|
def image_to_base64(img):
|
|
"""Convert PIL image to base64 string."""
|
|
buffer = io.BytesIO()
|
|
img.save(buffer, format='PNG')
|
|
return base64.b64encode(buffer.getvalue()).decode('utf-8')
|
|
|
|
|
|
def main():
|
|
"""Demonstrate the ImageLoader node functionality."""
|
|
print("ComfyUI Universal Image Loader - Example Usage")
|
|
print("=" * 50)
|
|
|
|
# Create an instance of the node
|
|
loader = ImageLoader()
|
|
|
|
# Create a test image
|
|
test_img = create_test_image()
|
|
test_base64 = image_to_base64(test_img)
|
|
test_data_url = f"data:image/png;base64,{test_base64}"
|
|
|
|
print("1. Testing Base64 Input")
|
|
print("-" * 20)
|
|
|
|
# Test with base64 string
|
|
image_tensor, mask_tensor = loader.load_image(base64=test_base64)
|
|
print(f"Image tensor shape: {image_tensor.shape}")
|
|
print(f"Mask tensor shape: {mask_tensor.shape}")
|
|
print(f"Image value range: {image_tensor.min():.3f} - {image_tensor.max():.3f}")
|
|
print()
|
|
|
|
print("2. Testing Data URL Input")
|
|
print("-" * 20)
|
|
|
|
# Test with data URL
|
|
image_tensor, mask_tensor = loader.load_image(base64=test_data_url)
|
|
print(f"Image tensor shape: {image_tensor.shape}")
|
|
print(f"Mask tensor shape: {mask_tensor.shape}")
|
|
print()
|
|
|
|
print("3. Testing Input Precedence")
|
|
print("-" * 20)
|
|
|
|
# Test precedence: pasted_base64 should take priority
|
|
image_tensor, mask_tensor = loader.load_image(
|
|
filepath="nonexistent.png",
|
|
base64="invalid_base64",
|
|
pasted_base64=test_data_url
|
|
)
|
|
print(f"With multiple inputs, pasted_base64 was used:")
|
|
print(f"Image tensor shape: {image_tensor.shape}")
|
|
print(f"Success: {image_tensor.shape != (1, 1, 1, 3)}")
|
|
print()
|
|
|
|
print("4. Testing Empty Input Handling")
|
|
print("-" * 20)
|
|
|
|
# Test with no inputs
|
|
image_tensor, mask_tensor = loader.load_image()
|
|
print(f"Empty input result:")
|
|
print(f"Image tensor shape: {image_tensor.shape}")
|
|
print(f"Mask tensor shape: {mask_tensor.shape}")
|
|
print(f"Is fallback: {image_tensor.shape == (1, 1, 1, 3)}")
|
|
print()
|
|
|
|
print("5. Testing Alpha Channel Handling")
|
|
print("-" * 20)
|
|
|
|
# Create RGBA image with alpha channel
|
|
rgba_img = Image.new('RGBA', (100, 100), (255, 0, 0, 128))
|
|
rgba_base64 = image_to_base64(rgba_img)
|
|
|
|
image_tensor, mask_tensor = loader.load_image(base64=rgba_base64)
|
|
print(f"RGBA image processing:")
|
|
print(f"Image tensor shape: {image_tensor.shape}")
|
|
print(f"Mask tensor shape: {mask_tensor.shape}")
|
|
print(f"Alpha values (should be ~0.5): {mask_tensor.mean():.3f}")
|
|
print()
|
|
|
|
print("6. Node Configuration Info")
|
|
print("-" * 20)
|
|
|
|
# Display node configuration
|
|
input_types = ImageLoader.INPUT_TYPES()
|
|
print(f"Node category: {ImageLoader.CATEGORY}")
|
|
print(f"Return types: {ImageLoader.RETURN_TYPES}")
|
|
print(f"Return names: {ImageLoader.RETURN_NAMES}")
|
|
print(f"Function name: {ImageLoader.FUNCTION}")
|
|
print(f"Optional inputs: {list(input_types['optional'].keys())}")
|
|
print()
|
|
|
|
print("Example completed successfully!")
|
|
print("The ImageLoader node is ready for use in ComfyUI workflows.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|