Fix linting
This commit is contained in:
@@ -1,6 +1,6 @@
|
||||
import torch
|
||||
from inspect import cleandoc
|
||||
from typing import Any, Union
|
||||
from typing import Any
|
||||
|
||||
try:
|
||||
from comfy.comfy_types.node_typing import IO, ComfyNodeABC
|
||||
|
||||
+12
-13
@@ -1,4 +1,3 @@
|
||||
import pytest
|
||||
import torch
|
||||
from src.basic_data_handling.tensor_nodes import (
|
||||
TensorCreate, TensorBinaryOp, TensorUnaryOp, TensorSlice,
|
||||
@@ -11,11 +10,11 @@ def test_tensor_create():
|
||||
result = node.create([1, 2, 3])
|
||||
assert isinstance(result[0], torch.Tensor)
|
||||
assert torch.equal(result[0], torch.tensor([1, 2, 3]))
|
||||
|
||||
|
||||
# From scalar
|
||||
result = node.create(5.0)
|
||||
assert torch.equal(result[0], torch.tensor(5.0))
|
||||
|
||||
|
||||
# From existing tensor
|
||||
t = torch.tensor([1.0, 2.0])
|
||||
result = node.create(t)
|
||||
@@ -25,7 +24,7 @@ def test_tensor_binary_op():
|
||||
node = TensorBinaryOp()
|
||||
a = torch.tensor([10.0, 20.0])
|
||||
b = torch.tensor([2.0, 4.0])
|
||||
|
||||
|
||||
assert torch.equal(node.operate(a, b, "add")[0], a + b)
|
||||
assert torch.equal(node.operate(a, b, "subtract")[0], a - b)
|
||||
assert torch.equal(node.operate(a, b, "multiply")[0], a * b)
|
||||
@@ -35,7 +34,7 @@ def test_tensor_binary_op():
|
||||
def test_tensor_unary_op():
|
||||
node = TensorUnaryOp()
|
||||
t = torch.tensor([-1.0, 0.0, 1.0])
|
||||
|
||||
|
||||
assert torch.equal(node.operate(t, "abs")[0], torch.tensor([1.0, 0.0, 1.0]))
|
||||
assert torch.equal(node.operate(t, "neg")[0], torch.tensor([1.0, 0.0, -1.0]))
|
||||
assert torch.equal(node.operate(torch.tensor([0.0]), "sin")[0], torch.tensor([0.0]))
|
||||
@@ -43,11 +42,11 @@ def test_tensor_unary_op():
|
||||
def test_tensor_slice():
|
||||
node = TensorSlice()
|
||||
t = torch.arange(10).reshape(2, 5) # [[0,1,2,3,4], [5,6,7,8,9]]
|
||||
|
||||
|
||||
# Simple slice
|
||||
result = node.slice_tensor(t, "0, 1:3")
|
||||
assert torch.equal(result[0], t[0, 1:3])
|
||||
|
||||
|
||||
# Ellipsis/Full slice
|
||||
result = node.slice_tensor(t, ":, 2")
|
||||
assert torch.equal(result[0], t[:, 2])
|
||||
@@ -55,17 +54,17 @@ def test_tensor_slice():
|
||||
def test_tensor_reshape():
|
||||
node = TensorReshape()
|
||||
t = torch.arange(6)
|
||||
|
||||
|
||||
result = node.reshape(t, "2, 3")
|
||||
assert result[0].shape == (2, 3)
|
||||
|
||||
|
||||
result = node.reshape(t, "-1, 2")
|
||||
assert result[0].shape == (3, 2)
|
||||
|
||||
def test_tensor_permute():
|
||||
node = TensorPermute()
|
||||
t = torch.randn(2, 3, 4)
|
||||
|
||||
|
||||
result = node.permute(t, "2, 0, 1")
|
||||
assert result[0].shape == (4, 2, 3)
|
||||
|
||||
@@ -73,11 +72,11 @@ def test_tensor_join():
|
||||
node = TensorJoin()
|
||||
t1 = torch.tensor([1, 2])
|
||||
t2 = torch.tensor([3, 4])
|
||||
|
||||
|
||||
# Cat
|
||||
result = node.join(t1, t2, 0, "concatenate")
|
||||
assert torch.equal(result[0], torch.tensor([1, 2, 3, 4]))
|
||||
|
||||
|
||||
# Stack
|
||||
result = node.join(t1, t2, 0, "stack")
|
||||
assert result[0].shape == (2, 2)
|
||||
@@ -86,7 +85,7 @@ def test_tensor_join():
|
||||
def test_tensor_info():
|
||||
node = TensorInfo()
|
||||
t = torch.zeros((2, 3), dtype=torch.float32)
|
||||
|
||||
|
||||
shape, dtype, device = node.get_info(t)
|
||||
assert shape == [2, 3]
|
||||
assert "float32" in dtype
|
||||
|
||||
Reference in New Issue
Block a user