From ebca09a6cb1a18deee1095068fc260c12ec6e3ec Mon Sep 17 00:00:00 2001
From: Gavroche <jlenoir2k@gmail.com>
Date: Fri, 10 Oct 2025 21:36:42 -0400
Subject: [PATCH] FIX: Add robust process_output handling and output validation

Fixed potential PIL compatibility issues:

1. Safe process_output handling:
   - Added try/catch around vae.process_output() calls
   - Validate output format after process_output
   - Fallback to original tensor if process_output fails or returns invalid format

2. Final output validation:
   - Validate tensor shape is 4D (B, C, H, W)
   - Validate channel count matches expected (3 for RGB)
   - Create valid fallback tensor if validation fails
   - Added detailed logging for debugging

3. Enhanced error recovery:
   - Multiple levels of fallback to prevent PIL errors
   - Comprehensive logging for troubleshooting
   - Graceful handling of edge cases

This should prevent the PIL error: 'Cannot handle this data type: (1, 1, 512), |u1'
by ensuring our node always returns valid (B, 3, H, W) float32 tensors.
---
 nodes.py | 32 +++++++++++++++++++++++++++++++-
 1 file changed, 31 insertions(+), 1 deletion(-)

diff --git a/nodes.py b/nodes.py
index 3150056..51e9649 100644
--- a/nodes.py
+++ b/nodes.py
@@ -575,7 +575,20 @@ class ROCMOptimizedVAEDecode:
                         out = vae.first_stage_model.decode(samples_processed)
                 
                 out = out.to(vae.output_device).float()
-                pixel_samples = vae.process_output(out)
+                
+                # CRITICAL FIX: Handle process_output safely to avoid PIL issues
+                if hasattr(vae, 'process_output'):
+                    try:
+                        pixel_samples = vae.process_output(out)
+                        # Validate the output format
+                        if len(pixel_samples.shape) != 4 or pixel_samples.shape[1] != 3:
+                            logging.warning(f"process_output returned invalid format: {pixel_samples.shape}, using original")
+                            pixel_samples = out
+                    except Exception as e:
+                        logging.warning(f"process_output failed: {e}, using original output")
+                        pixel_samples = out
+                else:
+                    pixel_samples = out
                 
                 # Handle WAN VAE output format - ensure correct shape and channels
                 if len(pixel_samples.shape) == 5:  # Video format (B, C, T, H, W)
@@ -635,6 +648,23 @@ class ROCMOptimizedVAEDecode:
             pixel_samples = pixel_samples.reshape(-1, pixel_samples.shape[-3], 
                                                 pixel_samples.shape[-2], pixel_samples.shape[-1])
         
+        # CRITICAL FIX: Final validation of output format
+        if len(pixel_samples.shape) != 4:
+            logging.error(f"Invalid output shape: {pixel_samples.shape}, expected 4D tensor")
+            # Create a valid fallback
+            B, C, H, W = samples_tensor.shape
+            expected_h, expected_w = H * upscale_factor, W * upscale_factor
+            pixel_samples = torch.zeros(B, expected_output_channels, expected_h, expected_w, dtype=torch.float32, device=samples_tensor.device)
+        
+        if pixel_samples.shape[1] != expected_output_channels:
+            logging.error(f"Invalid output channels: {pixel_samples.shape[1]}, expected {expected_output_channels}")
+            # Create a valid fallback
+            B, C, H, W = samples_tensor.shape
+            expected_h, expected_w = H * upscale_factor, W * upscale_factor
+            pixel_samples = torch.zeros(B, expected_output_channels, expected_h, expected_w, dtype=torch.float32, device=samples_tensor.device)
+        
+        logging.info(f"Final output validation: shape={pixel_samples.shape}, dtype={pixel_samples.dtype}")
+        
         # Move to output device (no movedim needed - match standard VAE decode)
         pixel_samples = pixel_samples.to(vae.output_device)