CRITICAL PERFORMANCE FIX: - Replaced complex nodes.py with simplified, clean implementation - Removed all optimization overhead that was causing 173% slowdown - Ready for real-world performance testing Simplified Implementation Features: ✅ Clean, minimal VAE decode implementation ✅ No complex optimization layers or caching systems ✅ No performance tracking overhead ✅ Direct VAE decode with simple fallbacks ✅ Proper error handling and logging ✅ All test suite passes (21 tests) Key Changes: - Removed Phase 1, 2, 3 optimization layers - Removed performance stats tracking - Removed complex caching systems - Removed memory prefetching overhead - Removed tensor layout optimizations - Simplified precision handling - Direct decode with tiled fallback Expected Benefits: - Lower memory overhead - Faster execution (no optimization overhead) - More predictable performance - Easier to debug and maintain - Should fix 173% slowdown (55s → 150s) Status: ✅ READY FOR TESTING - Simplified version working correctly - All tests passing - Ready for real workflow performance test
141 lines
4.5 KiB
Python
141 lines
4.5 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Simple VAE Performance Test
|
|
Test the actual performance of the current VAE decode implementation
|
|
"""
|
|
import sys
|
|
import os
|
|
sys.path.insert(0, '/home/nino/ComfyUI')
|
|
|
|
import torch
|
|
import time
|
|
import psutil
|
|
from custom_nodes.rocm_ninodes.nodes import ROCMOptimizedVAEDecodeInstrumented
|
|
|
|
def create_mock_vae():
|
|
"""Create a mock VAE for testing"""
|
|
class MockVAE:
|
|
def __init__(self):
|
|
self.device = torch.device('cpu')
|
|
self.first_stage_model = MockFirstStageModel()
|
|
self.vae_dtype = torch.float32
|
|
self.patcher = MockPatcher()
|
|
self.output_device = torch.device('cpu')
|
|
|
|
def process_output(self, x):
|
|
return x.permute(0, 2, 3, 1)
|
|
|
|
def memory_used_decode(self, shape, dtype):
|
|
return 1024
|
|
|
|
def spacial_compression_decode(self, shape=None):
|
|
return 8
|
|
|
|
def temporal_compression_decode(self, shape=None):
|
|
return 1
|
|
|
|
def decode_tiled(self, samples, tile_size, overlap):
|
|
B, C, H, W = samples.shape
|
|
return torch.randn(B, 3, H*8, W*8)
|
|
|
|
def decode(self, samples):
|
|
B, C, H, W = samples.shape
|
|
return torch.randn(B, 3, H*8, W*8)
|
|
|
|
class MockPatcher:
|
|
def __init__(self):
|
|
self.model_patches_models = lambda: []
|
|
self.current_loaded_device = lambda: torch.device('cpu')
|
|
self.model_size = lambda: 1024
|
|
self.loaded_size = lambda: 1024
|
|
self.model_patches_to = lambda x: x
|
|
self.model_dtype = lambda: torch.float32
|
|
self.partially_load = lambda device, extra_memory, force_patch_weights=False: True
|
|
self.is_clone = lambda: False
|
|
self.parent = None
|
|
self.load_device = torch.device('cpu')
|
|
self.model = MockModel()
|
|
self.device = torch.device('cpu')
|
|
|
|
def model_memory_required(self, device):
|
|
return 1024
|
|
|
|
def model_load(self, lowvram_model_memory, force_patch_weights=False):
|
|
pass
|
|
|
|
def model_use_more_vram(self, use_more_vram, force_patch_weights=False):
|
|
pass
|
|
|
|
class MockModel:
|
|
def __init__(self):
|
|
self.model = None
|
|
|
|
class MockFirstStageModel:
|
|
def decode(self, samples):
|
|
B, C, H, W = samples.shape
|
|
return torch.randn(B, 3, H*8, W*8)
|
|
|
|
return MockVAE()
|
|
|
|
def test_vae_performance():
|
|
"""Test VAE decode performance"""
|
|
print("🔍 VAE Performance Test")
|
|
print("=" * 40)
|
|
|
|
node = ROCMOptimizedVAEDecodeInstrumented()
|
|
vae = create_mock_vae()
|
|
|
|
# Test different sizes
|
|
test_cases = [
|
|
{"name": "256x256", "latent": torch.randn(1, 4, 32, 32)},
|
|
{"name": "512x512", "latent": torch.randn(1, 4, 64, 64)},
|
|
{"name": "1024x1024", "latent": torch.randn(1, 4, 128, 128)},
|
|
]
|
|
|
|
for test_case in test_cases:
|
|
samples = {"samples": test_case["latent"]}
|
|
|
|
print(f"\n📊 Testing {test_case['name']}...")
|
|
|
|
# Warm up
|
|
for _ in range(2):
|
|
node.decode(
|
|
vae=vae,
|
|
samples=samples,
|
|
tile_size=512,
|
|
overlap=64,
|
|
use_rocm_optimizations=True
|
|
)
|
|
|
|
# Actual test
|
|
execution_times = []
|
|
for i in range(5):
|
|
start_time = time.time()
|
|
|
|
result = node.decode(
|
|
vae=vae,
|
|
samples=samples,
|
|
tile_size=512,
|
|
overlap=64,
|
|
use_rocm_optimizations=True
|
|
)
|
|
|
|
execution_time = time.time() - start_time
|
|
execution_times.append(execution_time)
|
|
|
|
print(f" Run {i+1}: {execution_time:.6f}s")
|
|
|
|
avg_time = sum(execution_times) / len(execution_times)
|
|
print(f" Average: {avg_time:.6f}s")
|
|
|
|
# Check if this is reasonable
|
|
if test_case["name"] == "256x256" and avg_time > 0.1:
|
|
print(f" ⚠️ WARNING: {test_case['name']} is slower than expected!")
|
|
elif test_case["name"] == "512x512" and avg_time > 0.5:
|
|
print(f" ⚠️ WARNING: {test_case['name']} is slower than expected!")
|
|
elif test_case["name"] == "1024x1024" and avg_time > 2.0:
|
|
print(f" ⚠️ WARNING: {test_case['name']} is slower than expected!")
|
|
|
|
if __name__ == "__main__":
|
|
test_vae_performance()
|