lum3on-ComfyUI-FrameUtilitys/examples.py

"""
Test script for ComfyUI Frame Utility nodes.

Validates functionality and performance of FrameExtender and FrameReplacer nodes.
Run this script to ensure proper installation and functionality.
"""

import torch
import sys
import os

# Add current directory to path for imports
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

# Import with absolute imports for testing
import importlib.util
import sys

# Load modules directly
spec = importlib.util.spec_from_file_location("frame_extender", "frame_extender.py")
frame_extender_module = importlib.util.module_from_spec(spec)

spec = importlib.util.spec_from_file_location("frame_replacer", "frame_replacer.py")
frame_replacer_module = importlib.util.module_from_spec(spec)

spec = importlib.util.spec_from_file_location("utils", "utils.py")
utils_module = importlib.util.module_from_spec(spec)

# Execute modules
spec.loader.exec_module(utils_module)
sys.modules['utils'] = utils_module

spec = importlib.util.spec_from_file_location("frame_extender", "frame_extender.py")
frame_extender_module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(frame_extender_module)

spec = importlib.util.spec_from_file_location("frame_replacer", "frame_replacer.py")
frame_replacer_module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(frame_replacer_module)

# Get classes
FrameExtender = frame_extender_module.FrameExtender
FrameReplacer = frame_replacer_module.FrameReplacer
utils = utils_module

def create_test_video(frames: int, height: int = 64, width: int = 64, channels: int = 3) -> torch.Tensor:
    """Create a test video with gradient patterns for easy visual verification."""
    video = torch.zeros(frames, height, width, channels)
    
    for i in range(frames):
        # Create unique pattern for each frame
        frame_value = i / max(1, frames - 1)  # 0 to 1
        
        # Gradient pattern
        for h in range(height):
            for w in range(width):
                video[i, h, w, 0] = frame_value  # Red channel = frame number
                video[i, h, w, 1] = h / height   # Green channel = height gradient
                video[i, h, w, 2] = w / width    # Blue channel = width gradient
    
    return video

def test_frame_extender():
    """Test FrameExtender functionality."""
    print("Testing FrameExtender...")
    
    extender = FrameExtender()
    
    # Test basic extension
    source = create_test_video(5, 32, 32, 3)
    additional = create_test_video(3, 32, 32, 3)
    
    # Test end extension
    result = extender.extend_frames(source, additional, insert_position=-1, width=32, height=32)
    extended = result[0]
    assert extended.shape[0] == 8, f"Expected 8 frames, got {extended.shape[0]}"
    print("✓ End extension test passed")

    # Test beginning extension
    result = extender.extend_frames(source, additional, insert_position=0, width=32, height=32)
    extended = result[0]
    assert extended.shape[0] == 8, f"Expected 8 frames, got {extended.shape[0]}"
    print("✓ Beginning extension test passed")

    # Test middle extension
    result = extender.extend_frames(source, additional, insert_position=2, width=32, height=32)
    extended = result[0]
    assert extended.shape[0] == 8, f"Expected 8 frames, got {extended.shape[0]}"
    print("✓ Middle extension test passed")

    # Test with blending
    result = extender.extend_frames(source, additional, insert_position=-1, width=32, height=32, blend_frames=2)
    extended = result[0]
    assert extended.shape[0] == 10, f"Expected 10 frames (5+3+2 blend), got {extended.shape[0]}"
    print("✓ Blending test passed")

    # Test resolution change
    additional_large = create_test_video(3, 64, 64, 3)
    result = extender.extend_frames(source, additional_large, insert_position=-1, width=64, height=64)
    extended = result[0]
    assert extended.shape[1:3] == (64, 64), f"Expected (64, 64), got {extended.shape[1:3]}"
    print("✓ Resolution change test passed")
    
    print("FrameExtender tests completed successfully! ✅\n")

def test_frame_replacer():
    """Test FrameReplacer functionality."""
    print("Testing FrameReplacer...")
    
    replacer = FrameReplacer()
    
    # Test basic replacement
    source = create_test_video(10, 32, 32, 3)
    replacement = create_test_video(3, 32, 32, 3)
    
    # Test single frame replacement
    result = replacer.replace_frames(source, replacement, target_frame=2, replace_count=1, width=32, height=32)
    modified = result[0]
    assert modified.shape[0] == 10, f"Expected 10 frames, got {modified.shape[0]}"
    print("✓ Single frame replacement test passed")

    # Test multiple frame replacement
    result = replacer.replace_frames(source, replacement, target_frame=2, replace_count=3, width=32, height=32)
    modified = result[0]
    assert modified.shape[0] == 10, f"Expected 10 frames, got {modified.shape[0]}"
    print("✓ Multiple frame replacement test passed")

    # Test replacement with looping
    result = replacer.replace_frames(source, replacement, target_frame=2, replace_count=5, width=32, height=32, loop_replacement=True)
    modified = result[0]
    assert modified.shape[0] == 10, f"Expected 10 frames, got {modified.shape[0]}"
    print("✓ Looping replacement test passed")

    # Test edge blending
    result = replacer.replace_frames(source, replacement, target_frame=2, replace_count=3, width=32, height=32, blend_edges=True, blend_strength=0.5)
    modified = result[0]
    assert modified.shape[0] == 10, f"Expected 10 frames, got {modified.shape[0]}"
    print("✓ Edge blending test passed")

    # Test boundary conditions
    result = replacer.replace_frames(source, replacement, target_frame=0, replace_count=2, width=32, height=32)
    modified = result[0]
    assert modified.shape[0] == 10, f"Expected 10 frames, got {modified.shape[0]}"
    print("✓ Beginning replacement test passed")

    result = replacer.replace_frames(source, replacement, target_frame=8, replace_count=2, width=32, height=32)
    modified = result[0]
    assert modified.shape[0] == 10, f"Expected 10 frames, got {modified.shape[0]}"
    print("✓ End replacement test passed")
    
    print("FrameReplacer tests completed successfully! ✅\n")

def test_utils():
    """Test utility functions."""
    print("Testing utility functions...")
    
    # Test tensor validation
    valid_tensor = torch.randn(5, 32, 32, 3)
    try:
        utils.validate_image_tensor(valid_tensor, "test")
        print("✓ Tensor validation test passed")
    except Exception as e:
        print(f"✗ Tensor validation test failed: {e}")
        return
    
    # Test resize functionality
    source = torch.randn(3, 32, 32, 3)
    target = torch.randn(3, 64, 64, 3)
    resized = utils.resize_to_match(source, target)
    assert resized.shape[1:3] == target.shape[1:3], "Resize failed"
    print("✓ Resize test passed")
    
    # Test frame blending
    frame1 = torch.ones(1, 32, 32, 3) * 0.0
    frame2 = torch.ones(1, 32, 32, 3) * 1.0
    blended = utils.blend_frames(frame1, frame2, 0.5)
    expected_value = 0.5
    assert torch.allclose(blended, torch.ones_like(blended) * expected_value, atol=1e-6), "Blending failed"
    print("✓ Frame blending test passed")
    
    # Test safe indexing
    safe_idx = utils.safe_frame_index(10, -1, allow_negative=True)
    assert safe_idx == 9, f"Expected 9, got {safe_idx}"
    print("✓ Safe indexing test passed")
    
    print("Utility function tests completed successfully! ✅\n")

def test_performance():
    """Test performance with larger videos."""
    print("Testing performance with larger videos...")
    
    # Create larger test videos
    large_source = create_test_video(100, 128, 128, 3)
    large_additional = create_test_video(50, 128, 128, 3)
    
    extender = FrameExtender()
    replacer = FrameReplacer()
    
    # Test memory-efficient processing
    import time
    
    start_time = time.time()
    result = extender.extend_frames(large_source, large_additional, width=128, height=128, memory_efficient=True)
    extend_time = time.time() - start_time
    print(f"✓ Large video extension completed in {extend_time:.2f}s")

    start_time = time.time()
    result = replacer.replace_frames(large_source, large_additional[:20], target_frame=10, replace_count=20, width=128, height=128, memory_efficient=True)
    replace_time = time.time() - start_time
    print(f"✓ Large video replacement completed in {replace_time:.2f}s")
    
    print("Performance tests completed successfully! ✅\n")

def main():
    """Run all tests."""
    print("🎬 ComfyUI Frame Utility Node Tests 🎬\n")
    print("=" * 50)
    
    try:
        test_utils()
        test_frame_extender()
        test_frame_replacer()
        test_performance()
        
        print("=" * 50)
        print("🎉 All tests passed successfully! 🎉")
        print("\nNodes are ready for use in ComfyUI!")
        
    except Exception as e:
        print(f"❌ Test failed with error: {e}")
        import traceback
        traceback.print_exc()
        return False
    
    return True

if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)