71 lines
2.6 KiB
Python
71 lines
2.6 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Test script to verify that image dimension fixes work correctly.
|
|
This script simulates the ComfyUI image output validation.
|
|
"""
|
|
|
|
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
def test_image_dimensions():
|
|
"""Test that all image outputs have valid dimensions for ComfyUI"""
|
|
|
|
print("Testing image dimension fixes...")
|
|
|
|
# Simulate input image dimensions
|
|
luma_image = torch.randn(1, 1280, 728, 3)
|
|
print(f"Input image shape: {luma_image.shape}")
|
|
|
|
# Simulate problematic outputs that might cause the error
|
|
test_cases = [
|
|
("Valid 3D tensor", torch.randn(1280, 728, 3)),
|
|
("Valid 2D tensor", torch.randn(1280, 728)),
|
|
("Invalid 1x1x1 tensor", torch.randn(1, 1, 1)),
|
|
("Invalid single pixel", torch.randn(1, 1)),
|
|
("Wrong dimensions", torch.randn(512, 512, 3)),
|
|
]
|
|
|
|
def ensure_valid_image(img_tensor, reference_shape):
|
|
"""Ensure image tensor has valid dimensions and data type for ComfyUI"""
|
|
if img_tensor.dim() == 2:
|
|
img_tensor = img_tensor.unsqueeze(-1)
|
|
if img_tensor.shape[:2] != reference_shape[:2]:
|
|
img_tensor = torch.nn.functional.interpolate(
|
|
img_tensor.permute(2, 0, 1).unsqueeze(0),
|
|
size=(reference_shape[1], reference_shape[2]),
|
|
mode='nearest'
|
|
).squeeze(0).permute(1, 2, 0)
|
|
return img_tensor.clamp(0, 1).float()
|
|
|
|
for test_name, test_tensor in test_cases:
|
|
print(f"\nTesting: {test_name}")
|
|
print(f" Input shape: {test_tensor.shape}")
|
|
print(f" Input dtype: {test_tensor.dtype}")
|
|
|
|
try:
|
|
# Apply the fix
|
|
fixed_tensor = ensure_valid_image(test_tensor, luma_image.shape)
|
|
print(f" Fixed shape: {fixed_tensor.shape}")
|
|
print(f" Fixed dtype: {fixed_tensor.dtype}")
|
|
|
|
# Test PIL conversion (this is where the error occurred)
|
|
numpy_array = fixed_tensor.numpy()
|
|
if numpy_array.ndim == 3 and numpy_array.shape[2] == 1:
|
|
numpy_array = numpy_array.squeeze(2)
|
|
|
|
# Convert to uint8 for PIL
|
|
numpy_array = np.clip(numpy_array * 255, 0, 255).astype(np.uint8)
|
|
|
|
# Test PIL Image creation
|
|
pil_image = Image.fromarray(numpy_array)
|
|
print(f" PIL conversion: SUCCESS - {pil_image.size}")
|
|
|
|
except Exception as e:
|
|
print(f" PIL conversion: FAILED - {e}")
|
|
|
|
print("\n✅ All tests completed!")
|
|
|
|
if __name__ == "__main__":
|
|
test_image_dimensions()
|