iGavroche
740b347e92
- ROCmSamplerCustomAdvanced: new attention_backend input (auto/pytorch/ comfy_kitchen_int8/flash), installed on a clone of guider.model_options so the shared ModelPatcher stays clean and sparse-attention overrides still chain on top - Removed dead inputs: precision_mode, batch_optimization, flash_attention, test_steps; fix_empty_latent_channels now passes downscale ratios like stock; disable_pbar actually forwarded - ROCmVAEDecode: deleted unreachable VAE re-dtype machinery (first_stage_model .dtype is None for every shipped VAE); decode follows vae.vae_dtype and never mutates shared weights; fp16 accumulation now goes through apply_rocm_backend_settings (rdna3_5 pre-7.14 guard respected), same fix in ROCmGGUFLoader which also stopped forcing allow_tf32=False - Deleted ROCmSamplerCustomAdvancedBenchmark node (ROCmFluxBenchmark remains) - batch_cap field + *_BATCH_CAP constants, select_precision, dead tail of apply_rocm_backend_settings removed - conftest stubs extended (comfy.audio, model_patcher, ldm attention, folder_paths, interrupt hook): 184 passed, 24 skipped, 0 failed - Docs: node docs, README, Benchmarks, ARCHITECTURE refreshed to current stack (ROCm 7.16 / torch 2.15.0a0+rocm10.1.0)
433 lines
14 KiB
Python
433 lines
14 KiB
Python
"""
|
|
Unit tests for ROCMOptimizedVAEDecode
|
|
"""
|
|
import pytest
|
|
import torch
|
|
import sys
|
|
import os
|
|
|
|
# Add the custom nodes directory to path
|
|
sys.path.insert(0, '/home/nino/ComfyUI/custom_nodes/rocm-ninodes')
|
|
|
|
from rocm_nodes.core.vae import ROCmVAEDecode, _detect_vae_type
|
|
|
|
|
|
class ROCmVAEDecodeInstrumented(ROCmVAEDecode):
|
|
"""Instrumented subclass for testing with deterministic behavior"""
|
|
pass
|
|
|
|
|
|
class TestROCmVAEDecode:
|
|
"""Test suite for ROCmVAEDecode"""
|
|
|
|
def test_basic_decode(self, sample_vae, sample_latent):
|
|
"""Test basic VAE decode functionality"""
|
|
node = ROCmVAEDecodeInstrumented()
|
|
|
|
result = node.decode(
|
|
vae=sample_vae,
|
|
samples=sample_latent,
|
|
tile_size=512,
|
|
overlap=64,
|
|
use_rocm_optimizations=True
|
|
)
|
|
|
|
assert isinstance(result, tuple)
|
|
assert len(result) == 1
|
|
assert isinstance(result[0], torch.Tensor)
|
|
assert len(result[0].shape) == 4
|
|
|
|
def test_different_tile_sizes(self, sample_vae, sample_latent):
|
|
"""Test VAE decode with different tile sizes"""
|
|
node = ROCmVAEDecodeInstrumented()
|
|
|
|
tile_sizes = [256, 512, 768, 1024]
|
|
for tile_size in tile_sizes:
|
|
result = node.decode(
|
|
vae=sample_vae,
|
|
samples=sample_latent,
|
|
tile_size=tile_size,
|
|
overlap=64,
|
|
use_rocm_optimizations=True
|
|
)
|
|
|
|
assert isinstance(result, tuple)
|
|
assert len(result) == 1
|
|
assert isinstance(result[0], torch.Tensor)
|
|
|
|
def test_different_resolutions(self, sample_vae):
|
|
"""Test VAE decode with different resolutions"""
|
|
node = ROCmVAEDecodeInstrumented()
|
|
|
|
resolutions = [(256, 256), (512, 512), (1024, 1024)]
|
|
for w, h in resolutions:
|
|
latent = {
|
|
"samples": torch.randn(1, 4, h//8, w//8)
|
|
}
|
|
|
|
result = node.decode(
|
|
vae=sample_vae,
|
|
samples=latent,
|
|
tile_size=512,
|
|
overlap=64,
|
|
use_rocm_optimizations=True
|
|
)
|
|
|
|
assert isinstance(result, tuple)
|
|
assert len(result) == 1
|
|
assert isinstance(result[0], torch.Tensor)
|
|
assert len(result[0].shape) == 4
|
|
|
|
def test_batch_processing(self, sample_vae):
|
|
"""Test VAE decode with a batch of latents"""
|
|
node = ROCmVAEDecodeInstrumented()
|
|
|
|
batch_latent = {
|
|
"samples": torch.randn(2, 4, 32, 32)
|
|
}
|
|
|
|
result = node.decode(
|
|
vae=sample_vae,
|
|
samples=batch_latent,
|
|
tile_size=512,
|
|
overlap=64,
|
|
use_rocm_optimizations=True
|
|
)
|
|
|
|
assert isinstance(result, tuple)
|
|
assert len(result) == 1
|
|
assert isinstance(result[0], torch.Tensor)
|
|
assert result[0].shape[0] == 2
|
|
|
|
def test_memory_optimization(self, sample_vae, sample_latent):
|
|
"""Test VAE decode with memory optimization enabled/disabled"""
|
|
node = ROCmVAEDecodeInstrumented()
|
|
|
|
for memory_opt in [True, False]:
|
|
result = node.decode(
|
|
vae=sample_vae,
|
|
samples=sample_latent,
|
|
tile_size=512,
|
|
overlap=64,
|
|
use_rocm_optimizations=True,
|
|
)
|
|
|
|
assert isinstance(result, tuple)
|
|
assert len(result) == 1
|
|
assert isinstance(result[0], torch.Tensor)
|
|
|
|
def test_adaptive_tiling(self, sample_vae, sample_latent):
|
|
"""Test VAE decode with adaptive tiling enabled/disabled"""
|
|
node = ROCmVAEDecodeInstrumented()
|
|
|
|
for adaptive_tiling in [True, False]:
|
|
result = node.decode(
|
|
vae=sample_vae,
|
|
samples=sample_latent,
|
|
tile_size=512,
|
|
overlap=64,
|
|
use_rocm_optimizations=True,
|
|
)
|
|
|
|
assert isinstance(result, tuple)
|
|
assert len(result) == 1
|
|
assert isinstance(result[0], torch.Tensor)
|
|
|
|
def test_error_handling(self, sample_vae):
|
|
"""Test VAE decode error handling"""
|
|
node = ROCmVAEDecodeInstrumented()
|
|
|
|
invalid_samples = {"samples": "invalid"}
|
|
|
|
with pytest.raises(Exception):
|
|
node.decode(
|
|
vae=sample_vae,
|
|
samples=invalid_samples,
|
|
tile_size=512,
|
|
overlap=64,
|
|
use_rocm_optimizations=True
|
|
)
|
|
|
|
def test_performance_consistency(self, sample_vae, sample_latent):
|
|
"""Test that VAE decode produces consistent results"""
|
|
node = ROCmVAEDecodeInstrumented()
|
|
|
|
results = []
|
|
for _ in range(3):
|
|
result = node.decode(
|
|
vae=sample_vae,
|
|
samples=sample_latent,
|
|
tile_size=512,
|
|
overlap=64,
|
|
use_rocm_optimizations=True
|
|
)
|
|
results.append(result[0])
|
|
|
|
for result in results:
|
|
assert result.shape == results[0].shape
|
|
|
|
def test_large_image(self, sample_vae):
|
|
"""Test VAE decode with large image"""
|
|
node = ROCmVAEDecodeInstrumented()
|
|
|
|
large_latent = {
|
|
"samples": torch.randn(1, 4, 128, 128)
|
|
}
|
|
|
|
result = node.decode(
|
|
vae=sample_vae,
|
|
samples=large_latent,
|
|
tile_size=768,
|
|
overlap=96,
|
|
use_rocm_optimizations=True
|
|
)
|
|
|
|
assert isinstance(result, tuple)
|
|
assert len(result) == 1
|
|
assert isinstance(result[0], torch.Tensor)
|
|
assert len(result[0].shape) == 4
|
|
|
|
def test_video_decode(self, sample_vae, sample_video_latent):
|
|
"""Test VAE decode with video (5D) latent"""
|
|
node = ROCmVAEDecodeInstrumented()
|
|
|
|
result = node.decode(
|
|
vae=sample_vae,
|
|
samples=sample_video_latent,
|
|
tile_size=512,
|
|
overlap=64,
|
|
use_rocm_optimizations=True
|
|
)
|
|
|
|
assert isinstance(result, tuple)
|
|
assert len(result) == 1
|
|
assert isinstance(result[0], torch.Tensor)
|
|
# Video output should be 4D (flattened batch+frames)
|
|
assert len(result[0].shape) == 4
|
|
|
|
def test_compatibility_mode(self, sample_vae, sample_latent):
|
|
"""Test compatibility mode disables ROCm optimizations"""
|
|
node = ROCmVAEDecodeInstrumented()
|
|
|
|
result = node.decode(
|
|
vae=sample_vae,
|
|
samples=sample_latent,
|
|
tile_size=512,
|
|
overlap=64,
|
|
use_rocm_optimizations=True,
|
|
compatibility_mode=True
|
|
)
|
|
|
|
assert isinstance(result, tuple)
|
|
assert len(result) == 1
|
|
assert isinstance(result[0], torch.Tensor)
|
|
|
|
# ── NEW TESTS ──────────────────────────────────────────────────────────
|
|
|
|
def test_ltx_vae_detection(self, sample_ltx_vae):
|
|
"""Test that LTX Video VAE is correctly detected"""
|
|
vae_type = _detect_vae_type(sample_ltx_vae)
|
|
assert vae_type == "ltxv_vae", f"Expected ltxv_vae, got {vae_type}"
|
|
|
|
def test_wan_vae_detection(self, sample_wan_vae):
|
|
"""Test that WAN VAE is correctly detected"""
|
|
vae_type = _detect_vae_type(sample_wan_vae)
|
|
assert vae_type == "wan_vae", f"Expected wan_vae, got {vae_type}"
|
|
|
|
def test_pixel_space_vae_detection(self, sample_pixel_vae):
|
|
"""Test that pixel-space VAE is correctly detected"""
|
|
vae_type = _detect_vae_type(sample_pixel_vae)
|
|
assert vae_type == "pixel_space", f"Expected pixel_space, got {vae_type}"
|
|
|
|
def test_standard_vae_detection(self, sample_vae):
|
|
"""Test that standard VAE returns 'standard'"""
|
|
vae_type = _detect_vae_type(sample_vae)
|
|
assert vae_type == "standard", f"Expected standard, got {vae_type}"
|
|
|
|
def test_ltx_video_decode(self, sample_ltx_vae, sample_ltx_latent):
|
|
"""Test LTX video decode processes full video at once"""
|
|
node = ROCmVAEDecodeInstrumented()
|
|
|
|
result = node.decode(
|
|
vae=sample_ltx_vae,
|
|
samples=sample_ltx_latent,
|
|
tile_size=1024,
|
|
overlap=128,
|
|
use_rocm_optimizations=True,
|
|
)
|
|
|
|
assert isinstance(result, tuple)
|
|
assert len(result) == 1
|
|
assert isinstance(result[0], torch.Tensor)
|
|
assert len(result[0].shape) == 4
|
|
|
|
def test_pixel_space_decode(self, sample_pixel_vae, sample_pixel_latent):
|
|
"""Test pixel-space VAE decode skips all optimization logic"""
|
|
node = ROCmVAEDecodeInstrumented()
|
|
|
|
result = node.decode(
|
|
vae=sample_pixel_vae,
|
|
samples=sample_pixel_latent,
|
|
tile_size=512,
|
|
overlap=64,
|
|
use_rocm_optimizations=True,
|
|
)
|
|
|
|
assert isinstance(result, tuple)
|
|
assert len(result) == 1
|
|
assert isinstance(result[0], torch.Tensor)
|
|
assert len(result[0].shape) == 4
|
|
# Pixel-space output should be 4D BHWC (passthrough preserves content)
|
|
|
|
def test_wan_video_decode(self, sample_wan_vae, sample_wan_latent):
|
|
"""Test WAN video decode processes full video at once"""
|
|
node = ROCmVAEDecodeInstrumented()
|
|
|
|
result = node.decode(
|
|
vae=sample_wan_vae,
|
|
samples=sample_wan_latent,
|
|
tile_size=768,
|
|
overlap=96,
|
|
use_rocm_optimizations=True,
|
|
)
|
|
|
|
assert isinstance(result, tuple)
|
|
assert len(result) == 1
|
|
assert isinstance(result[0], torch.Tensor)
|
|
assert len(result[0].shape) == 4
|
|
|
|
def test_high_compression_tiled_decode(self, sample_ltx_vae):
|
|
"""Test that tiled decode handles high compression (32x) correctly"""
|
|
node = ROCmVAEDecodeInstrumented()
|
|
|
|
large_ltx_latent = {
|
|
"samples": torch.randn(1, 128, 4, 16, 24)
|
|
}
|
|
|
|
result = node.decode(
|
|
vae=sample_ltx_vae,
|
|
samples=large_ltx_latent,
|
|
tile_size=1024,
|
|
overlap=128,
|
|
use_rocm_optimizations=True,
|
|
)
|
|
|
|
assert isinstance(result, tuple)
|
|
assert len(result) == 1
|
|
assert isinstance(result[0], torch.Tensor)
|
|
assert len(result[0].shape) == 4
|
|
|
|
def test_compatibility_with_quantized_ltx(self, sample_ltx_vae, sample_ltx_latent):
|
|
"""Test compatibility mode works with LTX VAE"""
|
|
node = ROCmVAEDecodeInstrumented()
|
|
|
|
result = node.decode(
|
|
vae=sample_ltx_vae,
|
|
samples=sample_ltx_latent,
|
|
tile_size=1024,
|
|
overlap=128,
|
|
use_rocm_optimizations=True,
|
|
compatibility_mode=True
|
|
)
|
|
|
|
assert isinstance(result, tuple)
|
|
assert len(result) == 1
|
|
assert isinstance(result[0], torch.Tensor)
|
|
|
|
def test_temporal_tiling_ltx(self, sample_ltx_vae):
|
|
"""Test temporal tiling for LTX long video decode"""
|
|
node = ROCmVAEDecodeInstrumented()
|
|
|
|
long_ltx_latent = {
|
|
"samples": torch.randn(1, 128, 48, 4, 6)
|
|
}
|
|
|
|
result = node.decode(
|
|
vae=sample_ltx_vae,
|
|
samples=long_ltx_latent,
|
|
tile_size=1024,
|
|
overlap=128,
|
|
use_rocm_optimizations=True,
|
|
enable_temporal_tiling=True,
|
|
temporal_chunk_size=16,
|
|
temporal_overlap=2,
|
|
last_frame_fix=False,
|
|
)
|
|
|
|
assert isinstance(result, tuple)
|
|
assert len(result) == 1
|
|
assert isinstance(result[0], torch.Tensor)
|
|
assert len(result[0].shape) == 4
|
|
assert result[0].shape[0] > 0
|
|
|
|
def test_temporal_tiling_with_last_frame_fix(self, sample_ltx_vae):
|
|
"""Test temporal tiling with last_frame_fix enabled"""
|
|
node = ROCmVAEDecodeInstrumented()
|
|
|
|
latent = {
|
|
"samples": torch.randn(1, 128, 16, 4, 6)
|
|
}
|
|
|
|
result = node.decode(
|
|
vae=sample_ltx_vae,
|
|
samples=latent,
|
|
tile_size=1024,
|
|
overlap=128,
|
|
use_rocm_optimizations=True,
|
|
enable_temporal_tiling=True,
|
|
temporal_chunk_size=8,
|
|
temporal_overlap=2,
|
|
last_frame_fix=True,
|
|
)
|
|
|
|
assert isinstance(result, tuple)
|
|
assert len(result) == 1
|
|
assert isinstance(result[0], torch.Tensor)
|
|
assert len(result[0].shape) == 4
|
|
assert result[0].shape[0] > 0
|
|
|
|
def test_temporal_tiling_frame_count(self, sample_ltx_vae):
|
|
"""Test that temporal tiling with blending preserves all frames (no loss)"""
|
|
node = ROCmVAEDecodeInstrumented()
|
|
temporal_comp = 8
|
|
|
|
latent_frames = 48
|
|
chunk_size = 16
|
|
latent = {
|
|
"samples": torch.randn(1, 128, latent_frames, 4, 6)
|
|
}
|
|
|
|
result = node.decode(
|
|
vae=sample_ltx_vae,
|
|
samples=latent,
|
|
tile_size=1024,
|
|
overlap=128,
|
|
use_rocm_optimizations=True,
|
|
enable_temporal_tiling=True,
|
|
temporal_chunk_size=chunk_size,
|
|
temporal_overlap=2,
|
|
last_frame_fix=False,
|
|
)
|
|
|
|
total_frames = result[0].shape[0]
|
|
expected = 1 + (latent_frames - 1) * temporal_comp
|
|
assert total_frames == expected, (
|
|
f"Expected {expected} frames (full no-tile count), got {total_frames}"
|
|
)
|
|
|
|
def test_temporal_tiling_disabled_by_default(self, sample_ltx_vae, sample_ltx_latent):
|
|
"""Test that temporal tiling is disabled by default (uses direct decode)"""
|
|
node = ROCmVAEDecodeInstrumented()
|
|
|
|
result = node.decode(
|
|
vae=sample_ltx_vae,
|
|
samples=sample_ltx_latent,
|
|
tile_size=1024,
|
|
overlap=128,
|
|
use_rocm_optimizations=True,
|
|
)
|
|
|
|
assert isinstance(result, tuple)
|
|
assert len(result) == 1
|
|
assert isinstance(result[0], torch.Tensor)
|