APZmedia-Comfyui-LightDirection-estimation/docs/Image-Dimension-Fix.md

Image Dimension Fix for ComfyUI

Problem

The error Cannot handle this data type: (1, 1, 1), |u1 occurred when ComfyUI tried to create a PIL Image from a tensor with invalid dimensions. This happened because:

1. Invalid Shape: Some image outputs had shape (1, 1, 1) instead of proper image dimensions
2. Wrong Data Type: The tensor had data type |u1 (unsigned 8-bit) which PIL couldn't handle
3. Dimension Mismatch: IRE shadow masks had different dimensions than the input image

Root Cause

The IRE shadow analysis was creating masks with dimensions that didn't match the input image, causing ComfyUI's PreviewImage node to fail when trying to convert tensors to PIL Images.

Solution

1. Dimension Validation

Added checks to ensure all image outputs match the input image dimensions:

# Ensure IRE masks match input image dimensions
shadow_mask_raw = ire_results['transition_analysis']['shadow_mask']
soft_shadow_mask_raw = ire_results['transition_analysis']['soft_shadow_mask']

# Resize to match input image if needed
if shadow_mask_raw.shape != luma_image.shape[:3]:
    shadow_mask_raw = torch.nn.functional.interpolate(
        shadow_mask_raw.unsqueeze(0).unsqueeze(0).float(),
        size=(luma_image.shape[1], luma_image.shape[2]),
        mode='nearest'
    ).squeeze(0).squeeze(0)

2. Universal Image Safety Function

Created a safety function that ensures all image outputs are valid:

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()

3. Applied to All Image Outputs

Applied the safety function to all image outputs:

# Apply safety checks to all image outputs
debug_mask = ensure_valid_image(debug_mask, luma_image.shape)
lit_normals_viz = ensure_valid_image(lit_normals_viz, luma_image.shape)
cluster_delta_chart = ensure_valid_image(cluster_delta_chart, luma_image.shape)
x_threshold_preview = ensure_valid_image(x_threshold_preview, luma_image.shape)
y_threshold_preview = ensure_valid_image(y_threshold_preview, luma_image.shape)
false_color_ire = ensure_valid_image(false_color_ire, luma_image.shape)
shadow_mask = ensure_valid_image(shadow_mask, luma_image.shape)
soft_shadow_mask = ensure_valid_image(soft_shadow_mask, luma_image.shape)

What This Fixes

Before (Problematic)

• IRE masks could have wrong dimensions: (1, 1, 1) or (512, 512) instead of (1280, 728)
• Data type issues with |u1 causing PIL conversion failures
• ComfyUI PreviewImage nodes failing with TypeError

After (Fixed)

• All image outputs have consistent dimensions matching input image
• All tensors are properly converted to float and clamped to [0, 1] range
• PIL Image conversion works reliably
• ComfyUI PreviewImage nodes display correctly

Technical Details

Dimension Handling

• 2D tensors: Automatically add channel dimension with unsqueeze(-1)
• Wrong size: Resize using torch.nn.functional.interpolate with nearest neighbor
• Data type: Convert to float and clamp to [0, 1] range

Interpolation Method

• Uses mode='nearest' to preserve binary mask values
• Maintains mask integrity during resizing
• Prevents blurring of shadow boundaries

Memory Efficiency

• Only resizes when necessary
• Uses in-place operations where possible
• Maintains tensor device consistency

Testing

The fix ensures that:

1. All image outputs have valid dimensions matching the input image
2. Data types are compatible with PIL Image creation
3. No dimension mismatches occur between different analysis components
4. ComfyUI PreviewImage nodes work without errors

Benefits

• Reliability: Eliminates ComfyUI image display errors
• Consistency: All outputs have uniform dimensions
• Robustness: Handles edge cases and dimension mismatches
• Performance: Efficient resizing only when needed

This fix ensures that the IRE shadow analysis integrates seamlessly with ComfyUI's image handling system! 🎯