52 lines
1.7 KiB
Python
52 lines
1.7 KiB
Python
import torch
|
|
import sys
|
|
import os
|
|
|
|
# Add paths
|
|
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))) # ComfyUI root
|
|
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "more_math", "Parser")))
|
|
|
|
import optical_flow_utils as ofu
|
|
from comfy.nested_tensor import NestedTensor
|
|
|
|
def test_permute_and_apply():
|
|
print("\n--- Testing Permute and Flow Apply ---")
|
|
|
|
# 1. Test Permute on NestedTensor
|
|
print("Testing NestedTensor.permute...")
|
|
nt = NestedTensor([torch.randn(1, 256, 256, 3) for _ in range(2)])
|
|
try:
|
|
# Reorder [B, H, W, C] to [B, C, H, W]
|
|
nt_permuted = nt.permute(0, 3, 1, 2)
|
|
print(f"Original shape: {nt.shape}")
|
|
print(f"Permuted shape: {nt_permuted.shape}")
|
|
if nt_permuted.shape == (1, 3, 256, 256):
|
|
print("NestedTensor.permute SUCCESS!")
|
|
else:
|
|
print(f"NestedTensor.permute FAILED: Got {nt_permuted.shape}")
|
|
except Exception as e:
|
|
print(f"NestedTensor.permute FAILED with error: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
|
|
# 2. Test apply_flow
|
|
print("\nTesting ofu.apply_flow...")
|
|
img = torch.randn(1, 128, 128, 3)
|
|
flow = torch.zeros(1, 128, 128, 2) # Zero flow should be identity
|
|
flow[..., 0] = 10.0 # Shift 10px right
|
|
|
|
try:
|
|
warped = ofu.apply_flow(img, flow)
|
|
print(f"Warped shape: {warped.shape}")
|
|
if warped.shape == img.shape:
|
|
print("ofu.apply_flow SUCCESS!")
|
|
else:
|
|
print(f"ofu.apply_flow FAILED: Got {warped.shape}")
|
|
except Exception as e:
|
|
print(f"ofu.apply_flow FAILED with error: {e}")
|
|
import traceback
|
|
traceback.print_exc()
|
|
|
|
if __name__ == "__main__":
|
|
test_permute_and_apply()
|