Gavroche
00bb3de87a
- Added test_vae_error_scenarios.py with 9 comprehensive test cases - Added test_vae_optimization.py for performance optimization testing - Added debug_vae_decode.py for debugging VAE decode issues - Added VAE_OPTIMIZATION_PLAN.md with detailed optimization strategy - Added test_data/ directory structure for organized test data - All tests pass and prevent regression of critical errors
100 lines
3.1 KiB
Python
Executable File
100 lines
3.1 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""
|
|
Debug script for VAE decode issues
|
|
Loads the saved debug data and tests VAE decode in isolation
|
|
"""
|
|
|
|
import pickle
|
|
import torch
|
|
import sys
|
|
import os
|
|
|
|
# Add ComfyUI to path
|
|
sys.path.append('/home/nino/ComfyUI')
|
|
|
|
def debug_vae_decode():
|
|
"""Debug VAE decode with saved data"""
|
|
|
|
# Load debug data - find the most recent file
|
|
import glob
|
|
debug_files = glob.glob('test_data/debug/wan_vae_input_debug_*.pkl')
|
|
if not debug_files:
|
|
print("No debug data found in test_data/debug/")
|
|
print("Run the WAN workflow first to generate test data")
|
|
return
|
|
|
|
# Use the most recent file
|
|
latest_file = max(debug_files, key=os.path.getctime)
|
|
print(f"Loading debug data from: {latest_file}")
|
|
|
|
try:
|
|
with open(latest_file, 'rb') as f:
|
|
debug_data = pickle.load(f)
|
|
print("Loaded debug data:")
|
|
for key, value in debug_data.items():
|
|
if key != 'chunk_reshaped_tensor': # Don't print the large tensor
|
|
print(f" {key}: {value}")
|
|
except Exception as e:
|
|
print(f"Error loading debug data: {e}")
|
|
return
|
|
|
|
# Create a mock tensor with the same shape
|
|
shape = debug_data['chunk_reshaped_shape']
|
|
dtype = debug_data['chunk_reshaped_dtype']
|
|
device = debug_data['chunk_reshaped_device']
|
|
|
|
print(f"\nCreating mock tensor:")
|
|
print(f" Shape: {shape}")
|
|
print(f" Dtype: {dtype}")
|
|
print(f" Device: {device}")
|
|
|
|
# Create mock tensor
|
|
mock_tensor = torch.randn(shape, dtype=dtype, device=device)
|
|
|
|
# Test different input formats
|
|
print(f"\nTesting different input formats:")
|
|
|
|
# Format 1: Direct tensor (what we were doing wrong)
|
|
print(f"1. Direct tensor: {type(mock_tensor)}")
|
|
print(f" Shape: {mock_tensor.shape}")
|
|
|
|
# Format 2: LATENT dict (what we're trying)
|
|
latent_dict = {"samples": mock_tensor}
|
|
print(f"2. LATENT dict: {type(latent_dict)}")
|
|
print(f" Keys: {latent_dict.keys()}")
|
|
print(f" samples type: {type(latent_dict['samples'])}")
|
|
print(f" samples shape: {latent_dict['samples'].shape}")
|
|
|
|
# Format 3: Check what ComfyUI VAE actually expects
|
|
print(f"\n3. Checking ComfyUI VAE expectations...")
|
|
|
|
# Try to import ComfyUI VAE
|
|
try:
|
|
import comfy.sd
|
|
print(" ComfyUI VAE imported successfully")
|
|
|
|
# Check VAE decode signature
|
|
print(" Checking VAE decode method signature...")
|
|
# This will help us understand what format the VAE expects
|
|
|
|
except ImportError as e:
|
|
print(f" Failed to import ComfyUI: {e}")
|
|
|
|
# Save the mock data for further testing
|
|
test_data = {
|
|
'mock_tensor': mock_tensor,
|
|
'latent_dict': latent_dict,
|
|
'debug_data': debug_data
|
|
}
|
|
|
|
# Save test data with timestamp
|
|
timestamp = int(time.time())
|
|
filename = f'test_data/debug/wan_vae_test_data_{timestamp}.pkl'
|
|
os.makedirs('test_data/debug', exist_ok=True)
|
|
with open(filename, 'wb') as f:
|
|
pickle.dump(test_data, f)
|
|
print(f"\nSaved test data to {filename}")
|
|
|
|
if __name__ == "__main__":
|
|
debug_vae_decode()
|