Fix linting

This commit is contained in:
StableLlama
2026-02-08 22:00:25 +01:00
parent 26616079b3
commit 2861936fa3
2 changed files with 13 additions and 14 deletions
+1 -1
View File
@@ -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
View File
@@ -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