nodes/HowToSD-ComfyUI-Data-Analysis
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Removed PyTorch wrapper nodes. They were moved to ComfyUI-Py-Wrapper extension.

Browse Source
This commit is contained in:
HowToSD
2025-03-10 15:54:51 -07:00
parent 5fee9ffb60
commit 60692bf9c2
338 changed files with 6 additions and 14298 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/pytorch_wrapper/pt_abs.py
-39
View File
@@ -1,39 +0,0 @@
from typing import Any, Dict
import torch
class PtAbs:
"""
Computes the absolute value of each element in a PyTorch tensor.
category: PyTorch wrapper - Math operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens_a": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens_a: torch.Tensor) -> tuple:
"""
Computes the absolute value of a PyTorch tensor element-wise.
Args:
tens_a (torch.Tensor): A PyTorch tensor of numerical type.
Returns:
tuple: A tuple containing the resultant tensor after applying the absolute function.
"""
return (torch.abs(tens_a),)
modules/pytorch_wrapper/pt_acos.py
-39
View File
@@ -1,39 +0,0 @@
from typing import Any, Dict
import torch
class PtAcos:
"""
Computes the arccosine (inverse cosine) of a PyTorch tensor element-wise.
category: PyTorch wrapper - Math operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens_a": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens_a: torch.Tensor) -> tuple:
"""
Computes the element-wise arccosine (inverse cosine) of a PyTorch tensor.
Args:
tens_a (torch.Tensor): A PyTorch tensor of numerical type. Values must be in the range [-1, 1].
Returns:
tuple: A tuple containing the resultant tensor after applying the arccosine function.
"""
return (torch.acos(tens_a),)
modules/pytorch_wrapper/pt_add.py
-41
View File
@@ -1,41 +0,0 @@
from typing import Any, Dict
import torch
class PtAdd:
"""
Adds two PyTorch tensors.
category: PyTorch wrapper - Arithmetic operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens_a": ("TENSOR", {}),
"tens_b": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens_a: torch.Tensor, tens_b: torch.Tensor) -> tuple:
"""
Adds two PyTorch tensors.
Args:
tens_a (torch.Tensor): A PyTorch tensor of any shape.
tens_b (torch.Tensor): A PyTorch tensor of shape that supports addition with `tens_a`.
Returns:
tuple: A tuple containing the resultant tensor.
"""
return (torch.add(tens_a, tens_b),)
modules/pytorch_wrapper/pt_arange.py
-62
View File
@@ -1,62 +0,0 @@
import os
import sys
from typing import Any, Dict, Tuple
import torch
from .utils import DTYPE_MAPPING
PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
sys.path.append(MODULE_ROOT)
from cda.utils import str_to_number, str_to_number_with_default
class PtArange:
"""
Creates a PyTorch tensor using `torch.arange` with the specified start, end, and step values.
The values are parsed from string inputs and converted to numeric types. This allows users to specify
tensor range values in text fields. The function also ensures proper dtype selection.
category: PyTorch wrapper - Tensor creation
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"start": ("STRING", {"default":"", "multiline": False}),
"end": ("STRING", {"default":"", "multiline": False}),
"step": ("STRING", {"default":"", "multiline": False}),
"data_type": (("float32", "float16", "bfloat16", "float64", "uint8", "int8", "int16", "int32", "int64"),)
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, start: str, end: str, step: str, data_type: str) -> Tuple[torch.Tensor]:
"""
Creates a PyTorch tensor using `torch.arange` with the given start, stop, and step values.
Args:
start (str): The starting value of the sequence. If empty, defaults to 0.
end (str): The end value of the sequence (exclusive).
step (str): The step size for generating the sequence. If empty, defaults to 1.
data_type (str): The desired PyTorch dtype for the resulting tensor.
Returns:
Tuple[torch.Tensor]: A tuple containing the generated PyTorch tensor.
"""
start_value = str_to_number_with_default(start, 0)
end_value = str_to_number(end)
step_value = str_to_number_with_default(step, 1)
tens = torch.arange(start_value, end_value, step_value, dtype=DTYPE_MAPPING[data_type])
return (tens,)
modules/pytorch_wrapper/pt_argmax.py
-59
View File
@@ -1,59 +0,0 @@
from typing import Any, Dict
import torch
from .utils import str_to_dim
class PtArgmax:
"""
Computes the indices of the maximum values of a PyTorch tensor along the specified dimension(s).
Specify the dimension in the `dim` field using an integer as shown below:
```
0
```
category: PyTorch wrapper - Reduction operation & Summary statistics
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"tens": ("TENSOR", {}),
"dim": ("STRING", {"multiline": True}),
"keepdim": ("BOOLEAN", {"default": False})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, dim: str, keepdim: bool) -> tuple:
"""
Computes the indices of the maximum values along the specified dimension(s).
Args:
tens (torch.Tensor): The input tensor.
dim (str): A string representation of the dimension(s) along which to compute argmax,
e.g., "0" or "".
keepdim (bool): Whether to retain reduced dimensions.
Returns:
tuple: A tuple containing the tensor with indices of maximum values.
"""
dim = str_to_dim(dim)
# Ensure keepdim is a boolean
if not isinstance(keepdim, bool):
raise TypeError(f"Invalid type for keepdim: {keepdim}. Must be a boolean.")
# Compute the argmax along the specified dimension(s)
tens = torch.argmax(tens, dim=dim, keepdim=keepdim)
return (tens,)
modules/pytorch_wrapper/pt_argmin.py
-59
View File
@@ -1,59 +0,0 @@
from typing import Any, Dict
import torch
from .utils import str_to_dim
class PtArgmin:
"""
Computes the indices of the minimum values of a PyTorch tensor along the specified dimension(s).
Specify the dimension in the `dim` field using an integer as shown below:
```
0
```
category: PyTorch wrapper - Reduction operation & Summary statistics
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"tens": ("TENSOR", {}),
"dim": ("STRING", {"multiline": True}),
"keepdim": ("BOOLEAN", {"default": False})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, dim: str, keepdim: bool) -> tuple:
"""
Computes the indices of the minimum values along the specified dimension(s).
Args:
tens (torch.Tensor): The input tensor.
dim (str): A string representation of the dimension(s) along which to compute argmin,
e.g., "0" or "".
keepdim (bool): Whether to retain reduced dimensions.
Returns:
tuple: A tuple containing the tensor with indices of minimum values.
"""
dim = str_to_dim(dim)
# Ensure keepdim is a boolean
if not isinstance(keepdim, bool):
raise TypeError(f"Invalid type for keepdim: {keepdim}. Must be a boolean.")
# Compute the argmin along the specified dimension(s)
tens = torch.argmin(tens, dim=dim, keepdim=keepdim)
return (tens,)
modules/pytorch_wrapper/pt_asin.py
-39
View File
@@ -1,39 +0,0 @@
from typing import Any, Dict
import torch
class PtAsin:
"""
Computes the arcsine (inverse sine) of a PyTorch tensor element-wise.
category: PyTorch wrapper - Math operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens_a": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens_a: torch.Tensor) -> tuple:
"""
Computes the element-wise arcsine (inverse sine) of a PyTorch tensor.
Args:
tens_a (torch.Tensor): A PyTorch tensor of numerical type. Values must be in the range [-1, 1].
Returns:
tuple: A tuple containing the resultant tensor after applying the arcsine function.
"""
return (torch.asin(tens_a),)
modules/pytorch_wrapper/pt_atan.py
-39
View File
@@ -1,39 +0,0 @@
from typing import Any, Dict
import torch
class PtAtan:
"""
Computes the arc tangent (inverse tangent) of a PyTorch tensor element-wise.
category: PyTorch wrapper - Math operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens_a": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens_a: torch.Tensor) -> tuple:
"""
Computes the element-wise arc tangent (inverse tangent) of a PyTorch tensor.
Args:
tens_a (torch.Tensor): A PyTorch tensor of numerical type.
Returns:
tuple: A tuple containing the resultant tensor after applying the arc tangent function.
"""
return (torch.atan(tens_a),)
modules/pytorch_wrapper/pt_bitwise_and.py
-41
View File
@@ -1,41 +0,0 @@
from typing import Any, Dict
import torch
class PtBitwiseAnd:
"""
Performs a bitwise AND operation on two PyTorch tensors element-wise.
category: PyTorch wrapper - Bitwise operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens_a": ("TENSOR", {}),
"tens_b": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens_a: torch.Tensor, tens_b: torch.Tensor) -> tuple:
"""
Performs a bitwise AND operation on two PyTorch tensors.
Args:
tens_a (torch.Tensor): A PyTorch tensor of integer type.
tens_b (torch.Tensor): A PyTorch tensor of the same shape and integer type as `tens_a`.
Returns:
tuple: A tuple containing the resultant tensor after applying bitwise AND.
"""
return (torch.bitwise_and(tens_a, tens_b),)
modules/pytorch_wrapper/pt_bitwise_left_shift.py
-42
View File
@@ -1,42 +0,0 @@
from typing import Any, Dict
import torch
class PtBitwiseLeftShift:
"""
Performs a bitwise left shift operation on two PyTorch tensors element-wise.
category: PyTorch wrapper - Bitwise operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens_a": ("TENSOR", {}),
"tens_b": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens_a: torch.Tensor, tens_b: torch.Tensor) -> tuple:
"""
Performs a bitwise left shift operation on two PyTorch tensors.
Args:
tens_a (torch.Tensor): A PyTorch tensor of integer type.
tens_b (torch.Tensor): A PyTorch tensor of the same shape and integer type as `tens_a`,
representing the number of positions to shift.
Returns:
tuple: A tuple containing the resultant tensor after applying bitwise left shift.
"""
return (torch.bitwise_left_shift(tens_a, tens_b),)
modules/pytorch_wrapper/pt_bitwise_not.py
-39
View File
@@ -1,39 +0,0 @@
from typing import Any, Dict
import torch
class PtBitwiseNot:
"""
Performs a bitwise NOT operation on a PyTorch tensor element-wise.
category: PyTorch wrapper - Bitwise operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens_a": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens_a: torch.Tensor) -> tuple:
"""
Performs a bitwise NOT operation on a PyTorch tensor.
Args:
tens_a (torch.Tensor): A PyTorch tensor of integer type.
Returns:
tuple: A tuple containing the resultant tensor after applying bitwise NOT.
"""
return (torch.bitwise_not(tens_a),)
modules/pytorch_wrapper/pt_bitwise_or.py
-41
View File
@@ -1,41 +0,0 @@
from typing import Any, Dict
import torch
class PtBitwiseOr:
"""
Performs a bitwise OR operation on two PyTorch tensors element-wise.
category: PyTorch wrapper - Bitwise operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens_a": ("TENSOR", {}),
"tens_b": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens_a: torch.Tensor, tens_b: torch.Tensor) -> tuple:
"""
Performs a bitwise OR operation on two PyTorch tensors.
Args:
tens_a (torch.Tensor): A PyTorch tensor of integer type.
tens_b (torch.Tensor): A PyTorch tensor of the same shape and integer type as `tens_a`.
Returns:
tuple: A tuple containing the resultant tensor after applying bitwise OR.
"""
return (torch.bitwise_or(tens_a, tens_b),)
modules/pytorch_wrapper/pt_bitwise_right_shift.py
-42
View File
@@ -1,42 +0,0 @@
from typing import Any, Dict
import torch
class PtBitwiseRightShift:
"""
Performs a bitwise right shift operation on two PyTorch tensors element-wise.
category: PyTorch wrapper - Bitwise operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens_a": ("TENSOR", {}),
"tens_b": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens_a: torch.Tensor, tens_b: torch.Tensor) -> tuple:
"""
Performs a bitwise right shift operation on two PyTorch tensors.
Args:
tens_a (torch.Tensor): A PyTorch tensor of integer type.
tens_b (torch.Tensor): A PyTorch tensor of the same shape and integer type as `tens_a`,
representing the number of positions to shift.
Returns:
tuple: A tuple containing the resultant tensor after applying bitwise right shift.
"""
return (torch.bitwise_right_shift(tens_a, tens_b),)
modules/pytorch_wrapper/pt_bitwise_xor.py
-41
View File
@@ -1,41 +0,0 @@
from typing import Any, Dict
import torch
class PtBitwiseXor:
"""
Performs a bitwise XOR operation on two PyTorch tensors element-wise.
category: PyTorch wrapper - Bitwise operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens_a": ("TENSOR", {}),
"tens_b": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens_a: torch.Tensor, tens_b: torch.Tensor) -> tuple:
"""
Performs a bitwise XOR operation on two PyTorch tensors.
Args:
tens_a (torch.Tensor): A PyTorch tensor of integer type.
tens_b (torch.Tensor): A PyTorch tensor of the same shape and integer type as `tens_a`.
Returns:
tuple: A tuple containing the resultant tensor after applying bitwise XOR.
"""
return (torch.bitwise_xor(tens_a, tens_b),)
modules/pytorch_wrapper/pt_bmm.py
-44
View File
@@ -1,44 +0,0 @@
from typing import Any, Dict
import torch
class PtBmm:
"""
Performs batched matrix multiplication on two 3D PyTorch tensors.
category: PyTorch wrapper - Matrix operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens_a": ("TENSOR", {}),
"tens_b": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens_a: torch.Tensor, tens_b: torch.Tensor) -> tuple:
"""
Performs batched matrix multiplication on two 3D PyTorch tensors.
Args:
tens_a (torch.Tensor): A 3D PyTorch tensor (batch, m, n).
tens_b (torch.Tensor): A 3D PyTorch tensor (batch, n, p) that supports matrix multiplication with `tens_a`.
Returns:
tuple: A tuple containing the resultant 3D tensor.
"""
if tens_a.dim() != 3 or tens_b.dim() != 3:
raise ValueError("torch.bmm() only supports batch matrix multiplication for 3D tensors.")
return (torch.bmm(tens_a, tens_b),)
modules/pytorch_wrapper/pt_bool_create.py
-51
View File
@@ -1,51 +0,0 @@
from typing import Any, Dict
import torch
import ast
class PtBoolCreate:
"""
Creates a PyTorch tensor of dtype bool from True or False values entered as a list in the text field.
For example, entering:
```
[True, False, False]
```
will create a 1D boolean tensor.
category: PyTorch wrapper - Tensor creation
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"data": ("STRING", {"multiline": True}),
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, data: str) -> tuple:
"""
Creates a PyTorch tensor with bool data type using True or False string values entered in the text field.
Args:
data (str): String value in the text field.
Returns:
tuple: A tuple containing the tensor.
"""
# Convert string to Python list
list_data = ast.literal_eval(data)
# Convert list to PyTorch tensor
tensor = torch.tensor(list_data, dtype=torch.bool)
return (tensor,)
modules/pytorch_wrapper/pt_cos.py
-39
View File
@@ -1,39 +0,0 @@
from typing import Any, Dict
import torch
class PtCos:
"""
Computes the cosine of a PyTorch tensor element-wise.
category: PyTorch wrapper - Math operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens_a": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens_a: torch.Tensor) -> tuple:
"""
Computes the element-wise cosine of a PyTorch tensor.
Args:
tens_a (torch.Tensor): A PyTorch tensor of numerical type, representing angles in radians.
Returns:
tuple: A tuple containing the resultant tensor after applying the cosine function.
"""
return (torch.cos(tens_a),)
modules/pytorch_wrapper/pt_cosh.py
-39
View File
@@ -1,39 +0,0 @@
from typing import Any, Dict
import torch
class PtCosh:
"""
Computes the hyperbolic cosine of a PyTorch tensor element-wise.
category: PyTorch wrapper - Math operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens_a": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens_a: torch.Tensor) -> tuple:
"""
Computes the element-wise hyperbolic cosine of a PyTorch tensor.
Args:
tens_a (torch.Tensor): A PyTorch tensor of numerical type.
Returns:
tuple: A tuple containing the resultant tensor after applying the hyperbolic cosine function.
"""
return (torch.cosh(tens_a),)
modules/pytorch_wrapper/pt_crop.py
-65
View File
@@ -1,65 +0,0 @@
from typing import Any, Dict
import torch
class PtCrop:
"""
Crops a PyTorch tensor to the specified size. The input tensor must have a shape of (c, h, w) or (b, c, h, w).
In ComfyUI Analysis, the tensor data type is "TENSOR," while in ComfyUI, the type "IMAGE" is used. To pass image data from ComfyUI to a ComfyUI Analysis node (e.g., this interpolation node), first use the "Pt From Image" node to convert it to a tensor, then transpose the axes from (b, h, w, c) to (b, c, h, w) before passing it to this node.
This transposition can be performed using the "Pt Permute" node with (0, 2, 3, 1). To convert the output back to an image, apply "Pt Permute" with (0, 3, 1, 2) and then use the "Pt To Image" node.
category: PyTorch wrapper - Image processing
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens": ("TENSOR", {}),
"height": ("INT", {"default":0, "min":1, "max": 16384}),
"width": ("INT", {"default":0, "min":1, "max": 16384})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, height: int, width: int) -> tuple:
"""
Crops a PyTorch tensor.
Args:
tens (torch.Tensor): A PyTorch tensor with the shape of (c, h, w) or (b, c, h, w).
height (int): New image height.
width (int): New image width.
Returns:
tuple: A tuple containing a transformed tensor`.
"""
dim = tens.dim()
if dim == 3: # c, h, w
image_height = tens.size(1)
image_width = tens.size(2)
elif dim == 4: # b, c, h, w
image_height = tens.size(2)
image_width = tens.size(3)
else:
raise ValueError("Only rank 3 or 4 tensors are supported for crop.")
if height > image_height or width > image_width:
raise ValueError("Specified crop size exceeds tensor dimensions.")
# Compute starting position for cropping
top = max(int((height - image_height) / 2), 0)
left = max(int((width - image_width) / 2), 0)
tens2 = tens[..., top:top+height, left:left+width]
return (tens2,)
modules/pytorch_wrapper/pt_data_loader.py
-81
View File
@@ -1,81 +0,0 @@
import inspect
from typing import Any, Dict
import ast
import torch
from torch.utils.data import DataLoader
class PtDataLoader:
"""
Loads data from a dataset node and creates a PyTorch DataLoader.
**dataset**
Connect a PtvDataset node that provides the dataset.
**batch size**
Specifies the number of samples per batch.
**shuffle**
If True, shuffles the dataset before each epoch.
category: PyTorch wrapper - Training
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types required for data loading.
Returns:
Dict[str, Any]: A dictionary specifying required and optional input types.
"""
return {
"required": {
"dataset": ("PTVDATASET", {}),
"batch_size": ("INT", {"default": 1, "min": 1, "max": 1e6}),
"shuffle": ("BOOLEAN", {"default": False}),
},
"optional": {
"parameters": ("STRING", {"default": "", "multiline": True})
}
}
RETURN_TYPES: tuple = ("PTDATALOADER",)
FUNCTION: str = "f"
CATEGORY: str = "Training"
@classmethod
def IS_CHANGED(cls, **kw):
return float("NaN")
def f(self,
dataset: Any, # One of the classes defined in torchvision.datasets
batch_size: int,
shuffle: bool,
parameters: str) -> tuple:
"""
Creates a PyTorch DataLoader for the given dataset.
Args:
dataset (torchvision.datasets.{MNIST, ...}): An instance of a dataset from torchvision.datasets.
batch_size (int): The number of samples per batch.
shuffle (bool): If True, shuffles the dataset before each epoch.
parameters (str): Additional parameters in Python dictionary format.
Returns:
tuple: A tuple containing the created DataLoader instance.
"""
with torch.inference_mode(False):
param_dict = ast.literal_eval(parameters)
# Get valid parameters for the DataLoader class
valid_params = inspect.signature(DataLoader).parameters
# Filter only parameters that match the dataset's constructor
filtered_params = {k: v for k, v in param_dict.items() if k in valid_params}
dl = DataLoader(dataset=dataset,
batch_size=batch_size,
shuffle=shuffle,
**filtered_params)
return (dl,)
modules/pytorch_wrapper/pt_div.py
-41
View File
@@ -1,41 +0,0 @@
from typing import Any, Dict
import torch
class PtDiv:
"""
Divides one PyTorch tensor by another element-wise.
category: PyTorch wrapper - Arithmetic operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens_a": ("TENSOR", {}),
"tens_b": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens_a: torch.Tensor, tens_b: torch.Tensor) -> tuple:
"""
Divides one PyTorch tensor by another element-wise.
Args:
tens_a (torch.Tensor): A PyTorch tensor of any shape.
tens_b (torch.Tensor): A PyTorch tensor of shape that supports element-wise division with `tens_a`.
Returns:
tuple: A tuple containing the resultant tensor.
"""
return (torch.div(tens_a, tens_b),)
modules/pytorch_wrapper/pt_einsum.py
-67
View File
@@ -1,67 +0,0 @@
from typing import Any, Dict
import torch
class PtEinsum:
"""
Performs Tensor operations specified in the Einstein summation equation.
Specify Einstein summation notation in the first socket.
Specify two tensors in tens_a, tens_b. If you specify the third tensor, specify in tens_c.
Otherwise, leave the socket unconnected.
Sublist format is not supported.
category: PyTorch wrapper - Matrix operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"equation": ("STRING", {"default": ""}),
"tens_a": ("TENSOR", {})
},
"optional":{
"tens_b": ("TENSOR", {}),
"tens_c": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self,
equation: str,
tens_a: torch.Tensor,
tens_b: torch.Tensor = None,
tens_c: torch.Tensor = None) -> tuple:
"""
Performs Tensor operations specified in the Einstein summation equation.
Args:
equation (str): Einstein summation equation.
tens_a (torch.Tensor): A PyTorch tensor.
tens_b (torch.Tensor): A PyTorch tensor.
tens_c (torch.Tensor): A PyTorch tensor.
Returns:
tuple: A tuple containing the resultant tensor.
"""
if tens_a is None:
raise ValueError("tens_a is None")
if tens_b is None and tens_c is None:
tens_ret_val = torch.einsum(equation, tens_a)
elif tens_c is None:
tens_ret_val = torch.einsum(equation, tens_a, tens_b)
else:
tens_ret_val = torch.einsum(equation, tens_a, tens_b, tens_c)
return (tens_ret_val,)
modules/pytorch_wrapper/pt_eq.py
-41
View File
@@ -1,41 +0,0 @@
from typing import Any, Dict
import torch
class PtEq:
"""
Tests whether two PyTorch tensors are equal element-wise.
category: PyTorch wrapper - Comparison operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens_a": ("TENSOR", {}),
"tens_b": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens_a: torch.Tensor, tens_b: torch.Tensor) -> tuple:
"""
Compares two PyTorch tensors for element-wise equality.
Args:
tens_a (torch.Tensor): A PyTorch tensor of any shape.
tens_b (torch.Tensor): A PyTorch tensor of the same shape as `tens_a`.
Returns:
tuple: A tuple containing a boolean tensor indicating element-wise equality.
"""
return (torch.eq(tens_a, tens_b),)
modules/pytorch_wrapper/pt_evaluate_classification_model.py
-71
View File
@@ -1,71 +0,0 @@
from typing import Any, Dict, Tuple
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
class PtEvaluateClassificationModel:
"""
Performs inference on test data and computes evaluation metrics.
category: PyTorch wrapper - Training
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"model": ("PTMODEL", {}),
"data_loader": ("PTDATALOADER", {}),
"use_gpu": ("BOOLEAN", {"default":False})
}
}
RETURN_NAMES: Tuple[str, ...] = ("accuracy", "precision", "recall", "f1", "num_samples")
RETURN_TYPES: Tuple[str, ...] = ("FLOAT", "FLOAT", "FLOAT", "FLOAT", "INT")
FUNCTION: str = "f"
CATEGORY: str = "Training"
def f(self, model: nn.Module, data_loader: torch.utils.data.DataLoader, use_gpu: bool) -> Tuple[float, float, float, float, int]:
"""
Performs inference on test data and computes evaluation metrics.
Args:
model (torch.nn.Module): The PyTorch model to evaluate.
data_loader (torch.utils.data.DataLoader): DataLoader for the test dataset.
use_gpu (bool): Whether to use GPU for inference.
Returns:
Tuple[float, float, float, float, int]: A tuple containing accuracy, precision, recall,
F1-score, and the number of samples.
"""
with torch.inference_mode():
if use_gpu:
model.to("cuda")
model.train(False) # TODO: Change to model.*e*v*a*l() once Comfy's security checker is fixed.
y_hat = []
y = []
for images, labels in data_loader:
if use_gpu:
images = images.to("cuda")
labels = labels.to("cuda")
outputs = model(images)
preds = outputs.argmax(dim=1)
y_hat.extend(preds.cpu().numpy())
y.extend(labels.cpu().numpy())
accuracy = accuracy_score(y, y_hat)
precision = precision_score(y, y_hat, average='macro')
recall = recall_score(y, y_hat, average='macro')
f1 = f1_score(y, y_hat, average='macro')
num_samples = len(y)
return (accuracy, precision, recall, f1, num_samples)
modules/pytorch_wrapper/pt_exp.py
-39
View File
@@ -1,39 +0,0 @@
from typing import Any, Dict
import torch
class PtExp:
"""
Performs an exponential operation on a PyTorch tensor element-wise.
category: PyTorch wrapper - Math operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens_a": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens_a: torch.Tensor) -> tuple:
"""
Computes the element-wise exponential of a PyTorch tensor.
Args:
tens_a (torch.Tensor): A PyTorch tensor of numerical type.
Returns:
tuple: A tuple containing the resultant tensor after applying the exponential function.
"""
return (torch.exp(tens_a),)
modules/pytorch_wrapper/pt_flatten.py
-40
View File
@@ -1,40 +0,0 @@
from typing import Any, Dict
import torch
class PtFlatten:
"""
Flattens a PyTorch tensor into a 1D tensor.
category: PyTorch wrapper - Transform
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens": ("TENSOR", {}),
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor) -> tuple:
"""
Flattens a PyTorch tensor into a 1D tensor.
Args:
tens (torch.Tensor): A PyTorch tensor of any shape.
Returns:
tuple: A tuple containing the flattened tensor.
"""
return (tens.flatten(),)
modules/pytorch_wrapper/pt_float_create.py
-47
View File
@@ -1,47 +0,0 @@
from typing import Any, Dict
import torch
import ast
class PtFloatCreate:
"""
Creates a PyTorch tensor with 32-bit floating point precision
using values entered in the text field.
category: PyTorch wrapper - Tensor creation
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"data": ("STRING", {"multiline": True}),
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, data: str) -> tuple:
"""
Creates a PyTorch tensor using values entered in the text field.
Args:
data (str): String value in the text field.
Returns:
tuple: A tuple containing the tensor.
"""
# Convert string to Python list
list_data = ast.literal_eval(data)
# Convert list to PyTorch tensor
tensor = torch.tensor(list_data, dtype=torch.float32)
return (tensor,)
modules/pytorch_wrapper/pt_floor_divide.py
-41
View File
@@ -1,41 +0,0 @@
from typing import Any, Dict
import torch
class PtFloorDiv:
"""
Performs element-wise floor division on two PyTorch tensors.
category: PyTorch wrapper - Arithmetic operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens_a": ("TENSOR", {}),
"tens_b": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens_a: torch.Tensor, tens_b: torch.Tensor) -> tuple:
"""
Performs element-wise floor division on two PyTorch tensors.
Args:
tens_a (torch.Tensor): A PyTorch tensor of any shape.
tens_b (torch.Tensor): A PyTorch tensor of shape that supports element-wise floor division with `tens_a`.
Returns:
tuple: A tuple containing the resultant tensor.
"""
return (torch.floor_divide(tens_a, tens_b),)
modules/pytorch_wrapper/pt_from_image.py
-40
View File
@@ -1,40 +0,0 @@
from typing import Any, Dict
import torch
class PtFromImage:
"""
Casts an Image tensor as a PyTorch tensor.
category: PyTorch wrapper - Image processing
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"image": ("IMAGE", {}),
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, image: torch.Tensor) -> tuple:
"""
Casts an Image tensor as a PyTorch tensor.
Args:
image (torch.Tensor): Image tensor
Returns:
tuple: A tuple containing the PyTorch tensor.
"""
return (image,)
modules/pytorch_wrapper/pt_from_image_transpose.py
-47
View File
@@ -1,47 +0,0 @@
from typing import Any, Dict
import torch
class PtFromImageTranspose:
"""
Casts an image tensor to a PyTorch tensor and transposes it from (H, W, C) to (C, H, W). For rank-4 inputs, the batch axis remains unchanged.
category: PyTorch wrapper - Image processing
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"image": ("IMAGE", {}),
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, image: torch.Tensor) -> tuple:
"""
Casts an Image tensor as a PyTorch tensor.
Args:
image (torch.Tensor): Image tensor
Returns:
tuple: A tuple containing the PyTorch tensor.
"""
if image.dim() == 3:
tens = image.permute(2, 0, 1)
elif image.dim() == 4:
tens = image.permute(0, 3, 1, 2)
else:
raise ValueError("Only rank 3 or 4 tensors are supported.")
return (tens,)
modules/pytorch_wrapper/pt_from_latent.py
-40
View File
@@ -1,40 +0,0 @@
from typing import Any, Dict
import torch
class PtFromLatent:
"""
Casts a latent tensor as a PyTorch tensor.
category: PyTorch wrapper - Tensor creation
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"latent": ("LATENT", {}),
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, latent: torch.Tensor) -> tuple:
"""
Casts a latent tensor as a PyTorch tensor.
Args:
latent (torch.Tensor): Latent tensor
Returns:
tuple: A tuple containing the PyTorch tensor.
"""
return (latent["samples"],)
modules/pytorch_wrapper/pt_from_numpy.py
-42
View File
@@ -1,42 +0,0 @@
from typing import Any, Dict, Tuple
import numpy as np
import torch
class PtFromNumpy:
"""
Converts a NumPy ndarray to a PyTorch tensor while preserving its data type.
category: PyTorch wrapper - Tensor creation
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"array": ("NDARRAY", {}),
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, array: np.ndarray) -> Tuple[torch.Tensor]:
"""
Converts a NumPy ndarray to a PyTorch tensor while preserving its data type.
Args:
array (np.ndarray): NumPy ndarray.
Returns:
tuple: A tuple containing the PyTorch tensor.
"""
tens = torch.tensor(array, dtype=torch.from_numpy(array).dtype)
return (tens,)
modules/pytorch_wrapper/pt_full.py
-66
View File
@@ -1,66 +0,0 @@
from typing import Any, Dict
import torch
import ast
from .utils import DTYPE_MAPPING
class PtFull:
"""
Creates a PyTorch tensor filled with a specified value using the size entered in the text field.
For example, if you want to create a 4D tensor of batch size=2, channel=3, height=96, width=32 filled with 7, enter:
[2,3,96,32] for size and 7 for value.
category: PyTorch wrapper - Tensor creation
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"value": ("STRING", {"multiline": False}),
"size": ("STRING", {"multiline": True}),
"data_type": (("float32", "float16", "bfloat16", "float64", "uint8", "int8", "int16", "int32", "int64", "bool"),)
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, value: str, size: str, data_type: str) -> tuple:
"""
Creates a PyTorch tensor filled with a specified value.
Args:
value (str): String representing the value to fill the tensor.
size (str): String representation of the tensor size.
data_type (str): Desired data type of the tensor.
Returns:
tuple: A tuple containing the PyTorch tensor.
"""
# Convert string to Python list
list_size = ast.literal_eval(size)
# Convert the value to the appropriate type
dtype = DTYPE_MAPPING[data_type]
if dtype in {torch.float32, torch.float16, torch.bfloat16, torch.float64}:
fill_value = float(value)
elif dtype in {torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64}:
fill_value = int(value)
elif dtype == torch.bool:
fill_value = value.lower() in {"true", "1"}
else:
raise ValueError(f"Unsupported data type: {data_type}")
# Convert list to PyTorch Size
sz = torch.Size(list_size)
tens = torch.full(sz, fill_value, dtype=dtype)
return (tens,)
modules/pytorch_wrapper/pt_gather.py
-94
View File
@@ -1,94 +0,0 @@
from typing import Any, Dict
import torch
import ast
class PtGather:
"""
Generates a tensor based on the index tensor using PyTorch's `gather` function.
Specify the `dim` field using an integer to indicate the axis along which to gather elements.
Each input tensor element is selected based on the index tensor.
For example, consider the following input tensor:
```
[ [10, 20, 30, 40, 50],
[100, 200, 300, 400, 500]]
```
and the corresponding index tensor:
```
[ [0, 4, 3, 2, 0],
[4, 3, 0, 0, 0]]
```
If `dim = 0`, the function scans row-wise, treating each index value as a row index.
If `dim = 1`, it scans column-wise, interpreting each index value as a column index.
Let's analyze the case where `dim = 1`:
* Start at row 0, column 0.
* The index at this position is `index[0, 0] = 0`.
* PyTorch selects the value at position 0 from the array `[10, 20, 30, 40, 50]`, which is `10`.
* The result at this stage is:
```
[[10, ?, ?, ?, ?]
[?, ?, ?, ?, ?]]
```
* Moving to column 1: `index[0, 1] = 4`.
* The value at index 4 in `[10, 20, 30, 40, 50]` is `50`.
* The updated result is:
```
[[10, 50, ?, ?, ?]
[?, ?, ?, ?, ?]]
```
* Repeating this process for all elements constructs the final gathered tensor.
category: PyTorch wrapper - Indexing and Slicing Operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"tens": ("TENSOR", {}),
"dim": ("INT", {"default": 0, "min": -10, "max": 10}),
"index": ("STRING", {"multiline": True})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, dim: int, index: str) -> tuple:
"""
Gathers elements along the specified dimension using an index tensor.
Args:
tens (torch.Tensor): The input tensor.
dim (int): The dimension along which to gather values.
index (str): A string representation of the index tensor, provided in list format.
Returns:
tuple: A tuple containing the gathered tensor.
"""
# Parse the index string to a Python list
index = ast.literal_eval(index)
# Convert the parsed list to a tensor
index_tensor = torch.tensor(index, dtype=torch.long, device=tens.device)
# Perform the gather operation
gathered_tens = torch.gather(tens, dim=dim, index=index_tensor)
return (gathered_tens,)
modules/pytorch_wrapper/pt_ge.py
-41
View File
@@ -1,41 +0,0 @@
from typing import Any, Dict
import torch
class PtGe:
"""
Tests whether elements in the first PyTorch tensor are greater than or equal to the corresponding elements in the second tensor.
category: PyTorch wrapper - Comparison operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens_a": ("TENSOR", {}),
"tens_b": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens_a: torch.Tensor, tens_b: torch.Tensor) -> tuple:
"""
Compares two PyTorch tensors for element-wise greater-than-or-equal-to condition.
Args:
tens_a (torch.Tensor): A PyTorch tensor of any shape.
tens_b (torch.Tensor): A PyTorch tensor of the same shape as `tens_a`.
Returns:
tuple: A tuple containing a boolean tensor indicating where `tens_a >= tens_b`.
"""
return (torch.ge(tens_a, tens_b),)
modules/pytorch_wrapper/pt_gt.py
-41
View File
@@ -1,41 +0,0 @@
from typing import Any, Dict
import torch
class PtGt:
"""
Tests whether elements in the first PyTorch tensor are greater than the corresponding elements in the second tensor.
category: PyTorch wrapper - Comparison operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens_a": ("TENSOR", {}),
"tens_b": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens_a: torch.Tensor, tens_b: torch.Tensor) -> tuple:
"""
Compares two PyTorch tensors for element-wise greater-than condition.
Args:
tens_a (torch.Tensor): A PyTorch tensor of any shape.
tens_b (torch.Tensor): A PyTorch tensor of the same shape as `tens_a`.
Returns:
tuple: A tuple containing a boolean tensor indicating where `tens_a > tens_b`.
"""
return (torch.gt(tens_a, tens_b),)
modules/pytorch_wrapper/pt_index_select.py
-78
View File
@@ -1,78 +0,0 @@
from typing import Any, Dict
import torch
import ast
class PtIndexSelect:
"""
Extracts elements from the input tensor along a specified dimension using an index tensor.
`dim` specifies the dimension along which the selection occurs.
* For example, for 2D tensors:
** If `dim=0`, `index` selects rows.
** If `dim=1`, `index` selects columns (for 2D tensors).
`index` is a list of indices that will be converted into a tensor.
Example:
```
Input Tensor:
[[1, 2, 3],
[40, 50, 60],
[700, 800, 900]]
index = "[0, 2]"
dim = 0
Output Tensor:
[[1, 2, 3],
[700, 800, 900]]
```
category: PyTorch wrapper - Indexing and Slicing Operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"tens": ("TENSOR", {}),
"dim": ("INT", {"default": 0, "min": -10, "max": 10}),
"index": ("STRING", {"multiline": True})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, dim: int, index: str) -> tuple:
"""
Extracts elements from the input tensor along the specified dimension.
Args:
tens (torch.Tensor): The input tensor.
dim (int): The dimension along which elements are selected.
index (str): A string representation of a list of indices, e.g., "[0, 2]".
Returns:
tuple: A tuple containing the selected tensor.
"""
# Convert the string index to a Python list
index_list = ast.literal_eval(index)
if not isinstance(index_list, list):
raise ValueError("Index must be a list of integers.")
# Convert the index list to a tensor
index_tensor = torch.tensor(index_list, dtype=torch.long)
# Perform index selection
result = torch.index_select(tens, dim=dim, index=index_tensor)
return (result,)
modules/pytorch_wrapper/pt_int_create.py
-47
View File
@@ -1,47 +0,0 @@
from typing import Any, Dict
import torch
import ast
class PtIntCreate:
"""
Creates a PyTorch tensor with 32-bit integer
using values entered in the text field.
category: PyTorch wrapper - Tensor creation
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"data": ("STRING", {"multiline": True}),
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, data: str) -> tuple:
"""
Creates a PyTorch tensor using values entered in the text field.
Args:
data (str): String value in the text field.
Returns:
tuple: A tuple containing the tensor.
"""
# Convert string to Python list
list_data = ast.literal_eval(data)
# Convert list to PyTorch tensor
tensor = torch.tensor(list_data, dtype=torch.int32)
return (tensor,)
modules/pytorch_wrapper/pt_interpolate_by_scale_factor.py
-56
View File
@@ -1,56 +0,0 @@
from typing import Any, Dict
import torch
class PtInterpolateByScaleFactor:
"""
Resizes a PyTorch tensor using interpolation by scale factor. The input tensor must have a shape of (c, h, w) or (b, c, h, w).
In ComfyUI Analysis, the tensor data type is "TENSOR," while in ComfyUI, the type "IMAGE" is used. To pass image data from ComfyUI to a ComfyUI Analysis node (e.g., this interpolation node), first use the "Pt From Image" node to convert it to a tensor, then transpose the axes from (b, h, w, c) to (b, c, h, w) before passing it to this node.
This transposition can be performed using the "Pt Permute" node with (0, 2, 3, 1). To convert the output back to an image, apply "Pt Permute" with (0, 3, 1, 2) and then use the "Pt To Image" node.
category: PyTorch wrapper - Image processing
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens": ("TENSOR", {}),
"scale_factor": ("FLOAT", {"default":2.0, "min":0.01, "max": 100.0, "step":0.01}),
"mode": (("nearest-exact", "bilinear", "bicubic", "nearest"),)
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, scale_factor: float, mode: str) -> tuple:
"""
Resizes a PyTorch tensor using interpolation.
Args:
tens (torch.Tensor): A PyTorch tensor with the shape of (c, h, w) or (b, c, h, w).
scale_factor (float): Scale factor.
mode (string): Interpolation mode.
Returns:
tuple: A tuple containing a transformed tensor`.
"""
if tens.dim() == 3:
rank3 = True
tens_adjusted = torch.unsqueeze(tens, 0)
else:
rank3 = False
tens_adjusted = tens
tens2 = torch.nn.functional.interpolate(tens_adjusted, scale_factor=scale_factor, mode=mode)
if rank3:
tens2 = torch.squeeze(tens2, 0)
return (tens2,)
modules/pytorch_wrapper/pt_interpolate_to_size.py
-60
View File
@@ -1,60 +0,0 @@
from typing import Any, Dict
import torch
class PtInterpolateToSize:
"""
Resizes a PyTorch tensor using interpolation. The input tensor must have a shape of (c, h, w) or (b, c, h, w).
In ComfyUI Analysis, the tensor data type is "TENSOR," while in ComfyUI, the type "IMAGE" is used. To pass image data from ComfyUI to a ComfyUI Analysis node (e.g., this interpolation node), first use the "Pt From Image" node to convert it to a tensor, then transpose the axes from (b, h, w, c) to (b, c, h, w) before passing it to this node.
This transposition can be performed using the "Pt Permute" node with (0, 2, 3, 1). To convert the output back to an image, apply "Pt Permute" with (0, 3, 1, 2) and then use the "Pt To Image" node.
category: PyTorch wrapper - Image processing
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens": ("TENSOR", {}),
"height": ("INT", {"default":0, "min":1, "max": 16384}),
"width": ("INT", {"default":0, "min":1, "max": 16384}),
"mode": (("nearest-exact", "bilinear", "bicubic", "nearest"),)
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, height: int, width: int, mode: str) -> tuple:
"""
Resizes a PyTorch tensor using interpolation.
Args:
tens (torch.Tensor): A PyTorch tensor with the shape of (c, h, w) or (b, c, h, w).
height (int): New image height.
width (int): New image width.
mode (string): Interpolation mode.
Returns:
tuple: A tuple containing a transformed tensor`.
"""
if tens.dim() == 3:
rank3 = True
tens_adjusted = torch.unsqueeze(tens, 0)
else:
rank3 = False
tens_adjusted = tens
tens2 = torch.nn.functional.interpolate(tens_adjusted, size=(height, width), mode=mode)
if rank3:
tens2 = torch.squeeze(tens2, 0)
return (tens2,)
modules/pytorch_wrapper/pt_le.py
-41
View File
@@ -1,41 +0,0 @@
from typing import Any, Dict
import torch
class PtLe:
"""
Tests whether elements in the first PyTorch tensor are less than or equal to the corresponding elements in the second tensor.
category: PyTorch wrapper - Comparison operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens_a": ("TENSOR", {}),
"tens_b": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens_a: torch.Tensor, tens_b: torch.Tensor) -> tuple:
"""
Compares two PyTorch tensors for element-wise less-than-or-equal-to condition.
Args:
tens_a (torch.Tensor): A PyTorch tensor of any shape.
tens_b (torch.Tensor): A PyTorch tensor of the same shape as `tens_a`.
Returns:
tuple: A tuple containing a boolean tensor indicating where `tens_a <= tens_b`.
"""
return (torch.le(tens_a, tens_b),)
modules/pytorch_wrapper/pt_linspace.py
-62
View File
@@ -1,62 +0,0 @@
import os
import sys
from typing import Any, Dict, Tuple
import torch
from .utils import DTYPE_MAPPING
PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", ".."))
MODULE_ROOT = os.path.join(PROJECT_ROOT, "modules")
sys.path.append(MODULE_ROOT)
from cda.utils import str_to_number, str_to_number_with_default
class PtLinspace:
"""
Creates a PyTorch tensor using `torch.linspace` with the specified start, end, and steps values.
The start and end values are parsed from string inputs and converted to a float.
The steps value is converted to an int.
category: PyTorch wrapper - Tensor creation
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"start": ("STRING", {"default":"", "multiline": False}),
"end": ("STRING", {"default":"", "multiline": False}),
"steps": ("STRING", {"default":"", "multiline": False}),
"data_type": (("float32", "float16", "bfloat16", "float64", "uint8", "int8", "int16", "int32", "int64"),)
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, start: str, end: str, steps: str, data_type: str) -> Tuple[torch.Tensor]:
"""
Creates a PyTorch tensor using `torch.arange` with the given start, end, and step values.
Args:
start (str): The starting value of the sequence.
end (str): The end value of the sequence (inclusive).
steps (str): The number of steps
data_type (str): The desired PyTorch dtype for the resulting tensor.
Returns:
Tuple[torch.Tensor]: A tuple containing the generated PyTorch tensor.
"""
start_value = float(start)
end_value = float(end)
steps_value = int(steps)
tens = torch.linspace(start_value, end_value, steps_value, dtype=DTYPE_MAPPING[data_type])
return (tens,)
modules/pytorch_wrapper/pt_load_model.py
-41
View File
@@ -1,41 +0,0 @@
import os
from typing import Any, Dict
import torch
from .utils import get_model_full_path
class PtLoadModel:
"""
A wrapper class for saving a PyTorch model.
category: PyTorch wrapper - Training
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"model": ("PTMODEL", {}),
"model_path": ("STRING", {"default": "model.pkl"})
}
}
RETURN_NAMES: tuple = ("model",)
RETURN_TYPES: tuple = ("PTMODEL",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
@classmethod
def IS_CHANGED(cls, **kw):
return float("NaN")
def f(self, model:Any, model_path: str):
model_full_path = get_model_full_path(model_path)
sd = torch.load(model_full_path, weights_only=True)
model.load_state_dict(sd)
return (model,)
modules/pytorch_wrapper/pt_log.py
-39
View File
@@ -1,39 +0,0 @@
from typing import Any, Dict
import torch
class PtLog:
"""
Computes the natural logarithm (log base e) of a PyTorch tensor element-wise.
category: PyTorch wrapper - Math operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens_a": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens_a: torch.Tensor) -> tuple:
"""
Computes the element-wise natural logarithm (log base e) of a PyTorch tensor.
Args:
tens_a (torch.Tensor): A PyTorch tensor of numerical type. Must contain only positive values.
Returns:
tuple: A tuple containing the resultant tensor after applying the natural logarithm.
"""
return (torch.log(tens_a),)
modules/pytorch_wrapper/pt_logical_and.py
-41
View File
@@ -1,41 +0,0 @@
from typing import Any, Dict
import torch
class PtLogicalAnd:
"""
Performs a logical AND operation on two PyTorch tensors element-wise.
category: PyTorch wrapper - Logical operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens_a": ("TENSOR", {}),
"tens_b": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens_a: torch.Tensor, tens_b: torch.Tensor) -> tuple:
"""
Performs a logical AND operation on two PyTorch tensors.
Args:
tens_a (torch.Tensor): A PyTorch tensor of boolean or numerical type.
tens_b (torch.Tensor): A PyTorch tensor of the same shape and type as `tens_a`.
Returns:
tuple: A tuple containing a boolean tensor where `tens_a AND tens_b` is applied element-wise.
"""
return (torch.logical_and(tens_a, tens_b),)
modules/pytorch_wrapper/pt_logical_not.py
-39
View File
@@ -1,39 +0,0 @@
from typing import Any, Dict
import torch
class PtLogicalNot:
"""
Performs a logical NOT operation on a PyTorch tensor element-wise.
category: PyTorch wrapper - Logical operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens_a": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens_a: torch.Tensor) -> tuple:
"""
Performs a logical NOT operation on a PyTorch tensor.
Args:
tens_a (torch.Tensor): A PyTorch tensor of boolean or numerical type.
Returns:
tuple: A tuple containing the resultant tensor after applying logical NOT.
"""
return (torch.logical_not(tens_a),)
modules/pytorch_wrapper/pt_logical_or.py
-41
View File
@@ -1,41 +0,0 @@
from typing import Any, Dict
import torch
class PtLogicalOr:
"""
Performs a logical OR operation on two PyTorch tensors element-wise.
category: PyTorch wrapper - Logical operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens_a": ("TENSOR", {}),
"tens_b": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens_a: torch.Tensor, tens_b: torch.Tensor) -> tuple:
"""
Performs a logical OR operation on two PyTorch tensors.
Args:
tens_a (torch.Tensor): A PyTorch tensor of boolean or numerical type.
tens_b (torch.Tensor): A PyTorch tensor of the same shape and type as `tens_a`.
Returns:
tuple: A tuple containing a boolean tensor where `tens_a OR tens_b` is applied element-wise.
"""
return (torch.logical_or(tens_a, tens_b),)
modules/pytorch_wrapper/pt_logical_xor.py
-41
View File
@@ -1,41 +0,0 @@
from typing import Any, Dict
import torch
class PtLogicalXor:
"""
Performs a logical XOR operation on two PyTorch tensors element-wise.
category: PyTorch wrapper - Logical operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens_a": ("TENSOR", {}),
"tens_b": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens_a: torch.Tensor, tens_b: torch.Tensor) -> tuple:
"""
Performs a logical XOR operation on two PyTorch tensors.
Args:
tens_a (torch.Tensor): A PyTorch tensor of boolean or numerical type.
tens_b (torch.Tensor): A PyTorch tensor of the same shape and type as `tens_a`.
Returns:
tuple: A tuple containing a boolean tensor where `tens_a XOR tens_b` is applied element-wise.
"""
return (torch.logical_xor(tens_a, tens_b),)
modules/pytorch_wrapper/pt_lt.py
-41
View File
@@ -1,41 +0,0 @@
from typing import Any, Dict
import torch
class PtLt:
"""
Tests whether elements in the first PyTorch tensor are less than the corresponding elements in the second tensor.
category: PyTorch wrapper - Comparison operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens_a": ("TENSOR", {}),
"tens_b": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens_a: torch.Tensor, tens_b: torch.Tensor) -> tuple:
"""
Compares two PyTorch tensors for element-wise less-than condition.
Args:
tens_a (torch.Tensor): A PyTorch tensor of any shape.
tens_b (torch.Tensor): A PyTorch tensor of the same shape as `tens_a`.
Returns:
tuple: A tuple containing a boolean tensor indicating where `tens_a < tens_b`.
"""
return (torch.lt(tens_a, tens_b),)
modules/pytorch_wrapper/pt_masked_select.py
-62
View File
@@ -1,62 +0,0 @@
from typing import Any, Dict
import torch
class PtMaskedSelect:
"""
Extracts elements from the input tensor whose corresponding value in `masked_tens` is `True`.
Example:
```
Input Tensor:
[[1, 2, 3],
[40, 50, 60],
[700, 800, 900]]
Masked Tensor:
[[True, False, False],
[False, True, False],
[False, False, True]]
Output Tensor:
[1, 50, 900]
```
category: PyTorch wrapper - Indexing and Slicing Operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"tens": ("TENSOR", {}),
"masked_tens": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, masked_tens: torch.Tensor) -> tuple:
"""
Extracts elements from `tens` where `masked_tens` is `True`.
Args:
tens (torch.Tensor): The input tensor.
masked_tens (torch.Tensor): A boolean tensor specifying which elements to select.
Returns:
tuple: A tuple containing the selected tensor.
"""
if masked_tens.dtype != torch.bool:
raise ValueError("masked_tens must be a boolean tensor (dtype=torch.bool)")
result = torch.masked_select(tens, masked_tens)
return (result,)
modules/pytorch_wrapper/pt_matmul.py
-41
View File
@@ -1,41 +0,0 @@
from typing import Any, Dict
import torch
class PtMatMul:
"""
Performs matrix multiplication on two PyTorch tensors.
category: PyTorch wrapper - Matrix operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens_a": ("TENSOR", {}),
"tens_b": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens_a: torch.Tensor, tens_b: torch.Tensor) -> tuple:
"""
Performs matrix multiplication on two PyTorch tensors.
Args:
tens_a (torch.Tensor): A PyTorch tensor of shape compatible for matrix multiplication.
tens_b (torch.Tensor): A PyTorch tensor of shape that supports matrix multiplication with `tens_a`.
Returns:
tuple: A tuple containing the resultant tensor.
"""
return (torch.matmul(tens_a, tens_b),)
modules/pytorch_wrapper/pt_max.py
-59
View File
@@ -1,59 +0,0 @@
from typing import Any, Dict
import torch
import ast
from .utils import str_to_dim
class PtMax:
"""
Computes the maximum values of a PyTorch tensor along the specified dimension(s).
Specify the dimension(s) in the `dim` field using an integer, a list, or a tuple, as shown below:
```
0, [0], or (1, 2)
```
category: PyTorch wrapper - Reduction operation & Summary statistics
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"tens": ("TENSOR", {}),
"dim": ("STRING", {"multiline": True}),
"keepdim": ("BOOLEAN", {"default": False})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, dim: str, keepdim: bool) -> tuple:
"""
Computes the maximum values along the specified dimension(s).
Args:
tens (torch.Tensor): The input tensor.
dim (str): A string representation of the dimension(s) along which to compute the maximum,
e.g., "0", "[0]", or "(1, 2)".
keepdim (bool): Whether to retain reduced dimensions.
Returns:
tuple: A tuple containing the tensor with maximum values.
"""
dim = str_to_dim(dim)
# Ensure keepdim is a boolean
if not isinstance(keepdim, bool):
raise TypeError(f"Invalid type for keepdim: {keepdim}. Must be a boolean.")
# Compute the max along the specified dimension(s)
tens = torch.amax(tens, dim=dim, keepdim=keepdim)
return (tens,)
modules/pytorch_wrapper/pt_mean.py
-61
View File
@@ -1,61 +0,0 @@
from typing import Any, Dict
import torch
from .utils import str_to_dim
class PtMean:
"""
Computes the mean of a PyTorch tensor along the specified dimension(s).
This node supports only floating-point or complex tensors as input.
Specify the dimension(s) in the `dim` field using an integer, a list, or a tuple, as shown below:
```
0, [0], or (1, 2)
```
category: PyTorch wrapper - Reduction operation & Summary statistics
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"tens": ("TENSOR", {}),
"dim": ("STRING", {"multiline": True}),
"keepdim": ("BOOLEAN", {"default": False})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, dim: str, keepdim: bool) -> tuple:
"""
Computes the mean along the specified dimension(s).
Args:
tens (torch.Tensor): The input tensor.
dim (str): A string representation of the dimension(s) along which to compute the mean,
e.g., "0", "[0]", or "(1, 2)".
keepdim (bool): Whether to retain reduced dimensions.
Returns:
tuple: A tuple containing the mean tensor.
"""
dim = str_to_dim(dim)
# Ensure keepdim is a boolean
if not isinstance(keepdim, bool):
raise TypeError(f"Invalid type for keepdim: {keepdim}. Must be a boolean.")
# Compute the mean along the specified dimension(s)
tens = torch.mean(tens, dim=dim, keepdim=keepdim)
return (tens,)
modules/pytorch_wrapper/pt_median.py
-66
View File
@@ -1,66 +0,0 @@
from typing import Any, Dict
import torch
from .utils import str_to_dim
class PtMedian:
"""
Computes the median of a PyTorch tensor along the specified dimension(s).
Specify the dimension(s) in the `dim` field using an integer as shown below:
```
0
```
or
```
1
```
Note that PtMedian calls torch.median(), which returns the lower of the two numbers when the true median falls between them.
category: PyTorch wrapper - Reduction operation & Summary statistics
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"tens": ("TENSOR", {}),
"dim": ("STRING", {"multiline": True}),
"keepdim": ("BOOLEAN", {"default": False})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, dim: str, keepdim: bool) -> tuple:
"""
Computes the median along the specified dimension(s).
Args:
tens (torch.Tensor): The input tensor.
dim (str): A string representation of the dimension(s) along which to compute the median,
e.g., "0", or "1".
keepdim (bool): Whether to retain reduced dimensions.
Returns:
tuple: A tuple containing the median tensor.
"""
dim = str_to_dim(dim)
if not isinstance(dim, int):
raise TypeError(f"Invalid type for dim: {dim}. Must be a integer.")
# Ensure keepdim is a boolean
if not isinstance(keepdim, bool):
raise TypeError(f"Invalid type for keepdim: {keepdim}. Must be a boolean.")
# Compute the median along the specified dimension(s)
tens, _ = torch.median(tens, dim=dim, keepdim=keepdim)
return (tens,)
modules/pytorch_wrapper/pt_min.py
-59
View File
@@ -1,59 +0,0 @@
from typing import Any, Dict
import torch
from .utils import str_to_dim
class PtMin:
"""
Computes the minimum values of a PyTorch tensor along the specified dimension(s).
Specify the dimension(s) in the `dim` field using an integer, a list, or a tuple, as shown below:
```
0, [0], or (1, 2)
```
category: PyTorch wrapper - Reduction operation & Summary statistics
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"tens": ("TENSOR", {}),
"dim": ("STRING", {"multiline": True}),
"keepdim": ("BOOLEAN", {"default": False})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, dim: str, keepdim: bool) -> tuple:
"""
Computes the minimum values along the specified dimension(s).
Args:
tens (torch.Tensor): The input tensor.
dim (str): A string representation of the dimension(s) along which to compute the minimum,
e.g., "0", "[0]", or "(1, 2)".
keepdim (bool): Whether to retain reduced dimensions.
Returns:
tuple: A tuple containing the tensor with minimum values.
"""
dim = str_to_dim(dim)
# Ensure keepdim is a boolean
if not isinstance(keepdim, bool):
raise TypeError(f"Invalid type for keepdim: {keepdim}. Must be a boolean.")
# Compute the min along the specified dimension(s)
tens = torch.amin(tens, dim=dim, keepdim=keepdim)
return (tens,)
modules/pytorch_wrapper/pt_mm.py
-44
View File
@@ -1,44 +0,0 @@
from typing import Any, Dict
import torch
class PtMm:
"""
Performs 2D matrix multiplication on two PyTorch tensors.
category: PyTorch wrapper - Matrix operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens_a": ("TENSOR", {}),
"tens_b": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens_a: torch.Tensor, tens_b: torch.Tensor) -> tuple:
"""
Performs 2D matrix multiplication on two PyTorch tensors.
Args:
tens_a (torch.Tensor): A 2D PyTorch tensor.
tens_b (torch.Tensor): A 2D PyTorch tensor that supports matrix multiplication with `tens_a`.
Returns:
tuple: A tuple containing the resultant tensor.
"""
if tens_a.dim() != 2 or tens_b.dim() != 2:
raise ValueError("torch.mm() only supports 2D tensors.")
return (torch.mm(tens_a, tens_b),)
modules/pytorch_wrapper/pt_mul.py
-41
View File
@@ -1,41 +0,0 @@
from typing import Any, Dict
import torch
class PtMul:
"""
Multiplies two PyTorch tensors element-wise.
category: PyTorch wrapper - Arithmetic operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens_a": ("TENSOR", {}),
"tens_b": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens_a: torch.Tensor, tens_b: torch.Tensor) -> tuple:
"""
Multiplies two PyTorch tensors element-wise.
Args:
tens_a (torch.Tensor): A PyTorch tensor of any shape.
tens_b (torch.Tensor): A PyTorch tensor of shape that supports element-wise multiplication with `tens_a`.
Returns:
tuple: A tuple containing the resultant tensor.
"""
return (torch.mul(tens_a, tens_b),)
modules/pytorch_wrapper/pt_ne.py
-41
View File
@@ -1,41 +0,0 @@
from typing import Any, Dict
import torch
class PtNe:
"""
Tests whether two PyTorch tensors are not equal element-wise.
category: PyTorch wrapper - Comparison operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens_a": ("TENSOR", {}),
"tens_b": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens_a: torch.Tensor, tens_b: torch.Tensor) -> tuple:
"""
Compares two PyTorch tensors for element-wise inequality.
Args:
tens_a (torch.Tensor): A PyTorch tensor of any shape.
tens_b (torch.Tensor): A PyTorch tensor of the same shape as `tens_a`.
Returns:
tuple: A tuple containing a boolean tensor indicating element-wise inequality.
"""
return (torch.ne(tens_a, tens_b),)
modules/pytorch_wrapper/pt_neg.py
-39
View File
@@ -1,39 +0,0 @@
from typing import Any, Dict
import torch
class PtNeg:
"""
Computes the negation of each element in a PyTorch tensor.
category: PyTorch wrapper - Math operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens_a": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens_a: torch.Tensor) -> tuple:
"""
Computes the negation of a PyTorch tensor element-wise.
Args:
tens_a (torch.Tensor): A PyTorch tensor of numerical type.
Returns:
tuple: A tuple containing the resultant tensor after applying negation.
"""
return (torch.neg(tens_a),)
modules/pytorch_wrapper/pt_ones.py
-53
View File
@@ -1,53 +0,0 @@
from typing import Any, Dict
import torch
import ast
from .utils import DTYPE_MAPPING
class PtOnes:
"""
Creates a PyTorch tensor of ones using the size entered in the text field.
For example, if you want to create a 4D tensor of batch size=2, channel=3, height=96, width=32, enter:
[2,3,96,32]
category: PyTorch wrapper - Tensor creation
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"size": ("STRING", {"multiline": True}),
"data_type": (("float32", "float16", "bfloat16", "float64", "uint8", "int8", "int16", "int32", "int64", "bool"),)
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, size: str, data_type: str) -> tuple:
"""
Creates a PyTorch tensor of ones using the size entered in the text field.
Args:
size (str): String value in the text field.
data_type (str): Desired data type of the tensor.
Returns:
tuple: A tuple containing the PyTorch tensor.
"""
# Convert string to Python list
list_size = ast.literal_eval(size)
# Convert list to PyTorch Size
sz = torch.Size(list_size)
tens = torch.ones(sz, dtype=DTYPE_MAPPING[data_type])
return (tens,)
modules/pytorch_wrapper/pt_pad.py
-72
View File
@@ -1,72 +0,0 @@
from typing import Any, Dict
import torch
class PtPad:
"""
Pads a PyTorch tensor to the specified size. Padded area will be black. The input tensor must have a shape of (c, h, w) or (b, c, h, w).
In ComfyUI Analysis, the tensor data type is "TENSOR," while in ComfyUI, the type "IMAGE" is used. To pass image data from ComfyUI to a ComfyUI Analysis node (e.g., this interpolation node), first use the "Pt From Image" node to convert it to a tensor, then transpose the axes from (b, h, w, c) to (b, c, h, w) before passing it to this node.
This transposition can be performed using the "Pt Permute" node with (0, 2, 3, 1). To convert the output back to an image, apply "Pt Permute" with (0, 3, 1, 2) and then use the "Pt To Image" node.
category: PyTorch wrapper - Image processing
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens": ("TENSOR", {}),
"height": ("INT", {"default":0, "min":1, "max": 16384}),
"width": ("INT", {"default":0, "min":1, "max": 16384})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, height: int, width: int) -> tuple:
"""
Pads a PyTorch tensor.
Args:
tens (torch.Tensor): A PyTorch tensor with the shape of (c, h, w) or (b, c, h, w).
height (int): New image height.
width (int): New image width.
Returns:
tuple: A tuple containing a transformed tensor.
"""
dim = tens.dim()
if dim == 3: # c, h, w
image_height = tens.size(1)
image_width = tens.size(2)
elif dim == 4: # b, c, h, w
image_height = tens.size(2)
image_width = tens.size(3)
else:
raise ValueError("Only rank 3 or 4 tensors are supported for padding.")
if height < image_height or width < image_width:
raise ValueError("Specified dimensions must be greater than or equal to tensor dimensions.")
# Compute starting position for padding
top = max((height - image_height) // 2, 0)
left = max((width - image_width) // 2, 0)
if dim == 3:
tens2 = torch.zeros(tens.size(0), height, width,
dtype=tens.dtype, device=tens.device)
else:
tens2 = torch.zeros(tens.size(0), tens.size(1), height, width,
dtype=tens.dtype, device=tens.device)
tens2[..., top:top+image_height, left:left+image_width] = tens
return (tens2,)
modules/pytorch_wrapper/pt_permute.py
-60
View File
@@ -1,60 +0,0 @@
from typing import Any, Dict
import torch
import ast
class PtPermute:
"""
Permutes the dimensions of a PyTorch tensor according to the specified order.
For example, if a tensor has shape (2, 3, 96, 32) and `new_axes="[0, 3, 1, 2]"`,
the output will have shape (2, 32, 3, 96).
Specify the new axis order in the `new_axes` field using brackets, as shown below:
```
[0, 3, 1, 2]
```
category: PyTorch wrapper - Transform
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"tens": ("TENSOR", {}),
"new_axes": ("STRING", {"multiline": True})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, new_axes: str) -> tuple:
"""
Permutes the dimensions of the input tensor based on the specified order.
Args:
tens (torch.Tensor): The input tensor.
new_axes (str): A string representation of the desired dimension order,
e.g., "[0, 3, 1, 2]" or "(0, 3, 1, 2)".
Returns:
tuple: A tuple containing the permuted PyTorch tensor.
"""
# Convert string to a tuple of integers
permute_order = tuple(ast.literal_eval(new_axes))
# Ensure the parsed order is valid
if not isinstance(permute_order, tuple) or not all(isinstance(i, int) for i in permute_order):
raise ValueError(f"Invalid format for new_axes: {new_axes}. Must be a tuple or list of integers.")
# Apply permutation
tens = tens.permute(*permute_order)
return (tens,)
modules/pytorch_wrapper/pt_pow.py
-41
View File
@@ -1,41 +0,0 @@
from typing import Any, Dict
import torch
class PtPow:
"""
Raises one PyTorch tensor to the power of another element-wise.
category: PyTorch wrapper - Arithmetic operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens_a": ("TENSOR", {}),
"tens_b": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens_a: torch.Tensor, tens_b: torch.Tensor) -> tuple:
"""
Raises one PyTorch tensor to the power of another element-wise.
Args:
tens_a (torch.Tensor): A PyTorch tensor of any shape.
tens_b (torch.Tensor): A PyTorch tensor of shape that supports element-wise exponentiation with `tens_a`.
Returns:
tuple: A tuple containing the resultant tensor.
"""
return (torch.pow(tens_a, tens_b),)
modules/pytorch_wrapper/pt_predict_classification_model.py
-66
View File
@@ -1,66 +0,0 @@
from typing import Any, Dict, Tuple
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
class PtPredictClassificationModel:
"""
Performs inference on input data.
category: PyTorch wrapper - Training
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"model": ("PTMODEL", {}),
"inputs": ("TENSOR", {}),
"class_id_to_name_map": ("PYDICT", {}),
"use_gpu": ("BOOLEAN", {"default":False})
}
}
RETURN_NAMES: Tuple[str, ...] = ("class_name", "class_id", "probability")
RETURN_TYPES: Tuple[str, ...] = ("STRING", "INT", "FLOAT")
FUNCTION: str = "f"
CATEGORY: str = "Training"
def f(self, model: nn.Module,
inputs: torch.Tensor,
class_id_to_name_map: Dict,
use_gpu: bool) -> Tuple[str, int, float]:
"""
Performs inference on test data and computes evaluation metrics.
Args:
model (torch.nn.Module): The PyTorch model to evaluate.
inputs (torch.Tensor): Input image to predict the class for.
class_id_to_name_map (Dict): A Python dictionary contains the mapping of
class ID to class name.
use_gpu (bool): Whether to use GPU for inference.
Returns:
Tuple[str, int, float]: A tuple containing "class_name", "class_id", "probability".
"""
with torch.inference_mode():
if use_gpu:
model.to("cuda")
model.train(False) # TODO: Change to model.*e*v*a*l() once Comfy's security checker is fixed.
x = torch.unsqueeze(inputs, 0)
if use_gpu:
x = x.to("cuda")
outputs = model(x)
class_id = outputs.argmax(dim=1)[0].detach().cpu().item()
prob = F.softmax(outputs, dim=1)[0][class_id].detach().cpu().item()
class_name = class_id_to_name_map[class_id]
return (class_name, class_id, prob)
modules/pytorch_wrapper/pt_prod.py
-65
View File
@@ -1,65 +0,0 @@
from typing import Any, Dict
import torch
from .utils import str_to_dim
class PtProd:
"""
Computes the product of a PyTorch tensor along the specified dimension(s).
Specify the dimension(s) in the `dim` field using an integer as shown below:
```
0
```
or
```
1
```
category: PyTorch wrapper - Reduction operation & Summary statistics
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"tens": ("TENSOR", {}),
"dim": ("STRING", {"multiline": True}),
"keepdim": ("BOOLEAN", {"default": False})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, dim: str, keepdim: bool) -> tuple:
"""
Computes the product along the specified dimension(s).
Args:
tens (torch.Tensor): The input tensor.
dim (str): A string representation of the dimension(s) along which to compute the product,
e.g., "0" or "1".
keepdim (bool): Whether to retain reduced dimensions.
Returns:
tuple: A tuple containing the product tensor.
"""
dim = str_to_dim(dim)
if not isinstance(dim, int):
raise TypeError(f"Invalid type for dim: {dim}. Must be a integer.")
# Ensure keepdim is a boolean
if not isinstance(keepdim, bool):
raise TypeError(f"Invalid type for keepdim: {keepdim}. Must be a boolean.")
# Compute the product along the specified dimension(s)
tens = torch.prod(tens, dim=dim, keepdim=keepdim)
return (tens,)
modules/pytorch_wrapper/pt_rand.py
-57
View File
@@ -1,57 +0,0 @@
from typing import Any, Dict
import torch
import ast
from .utils import DTYPE_MAPPING
class PtRand:
"""
Creates a PyTorch tensor with values sampled from a uniform distribution
using the size entered in the text field.
For example, if you want to create a 4D tensor of batch size=2, channel=3, height=96, width=32, enter:
[2,3,96,32]
category: PyTorch wrapper - Tensor creation
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"size": ("STRING", {"multiline": True}),
"data_type": (("float32", "float16", "bfloat16", "float64"),)
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
@classmethod
def IS_CHANGED(cls, **kw):
return float("NaN")
def f(self, size: str, data_type: str) -> tuple:
"""
Creates a PyTorch tensor with values sampled from a uniform distribution.
Args:
size (str): String value in the text field.
data_type (str): Desired floating-point data type of the tensor.
Returns:
tuple: A tuple containing the PyTorch tensor.
"""
# Convert string to Python list
list_size = ast.literal_eval(size)
# Convert list to PyTorch Size
sz = torch.Size(list_size)
tens = torch.rand(sz, dtype=DTYPE_MAPPING[data_type])
return (tens,)
modules/pytorch_wrapper/pt_rand_int.py
-72
View File
@@ -1,72 +0,0 @@
import time
from typing import Any, Dict
import torch
import ast
import copy
from .utils import DTYPE_MAPPING
class PtRandInt:
"""
Creates a PyTorch tensor filled with random integers within a specified range using the size entered in the text field.
For example, if you want to create a 4D tensor of batch size=2, channel=3, height=96, width=32 with random integers in the range [0, 10),
enter [2,3,96,32] for size, 0 for min_value, and 10 for max_value.
category: PyTorch wrapper - Tensor creation
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"min_value": ("INT", {"default": 0, "min": -2**31, "max": 2**31}),
"max_value": ("INT", {"default": 1, "min": -2**31, "max": 2**31}),
"size": ("STRING", {"multiline": True}),
"data_type": (("uint8", "int8", "int16", "int32", "int64"),)
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
@classmethod
def IS_CHANGED(cls, **kw):
return float("NaN")
def __init__(self):
# Initialize a separate generator with a different seed
self.gen = torch.Generator()
self.gen.manual_seed(int(time.time() * 1000000) % (2**32)) # Unique seed
def f(self, min_value: int, max_value: int, size: str, data_type: str) -> tuple:
"""
Creates a PyTorch tensor filled with random integers in the range [min_value, max_value).
Args:
min_value (int): Minimum value (inclusive).
max_value (int): Maximum value (exclusive).
size (str): String representation of the tensor size.
data_type (str): Desired integer data type of the tensor.
Returns:
tuple: A tuple containing the PyTorch tensor.
"""
# Convert string to Python list
list_size = ast.literal_eval(size)
# Convert list to PyTorch Size
sz = torch.Size(list_size)
# Ensure max_value is greater than min_value
if max_value <= min_value:
raise ValueError("max_value must be greater than min_value.")
# Generate random integers within the specified range
tens = torch.randint(min_value, max_value, sz, dtype=DTYPE_MAPPING[data_type], generator=self.gen)
return (tens,)
modules/pytorch_wrapper/pt_randn.py
-57
View File
@@ -1,57 +0,0 @@
from typing import Any, Dict
import torch
import ast
from .utils import DTYPE_MAPPING
class PtRandn:
"""
Creates a PyTorch tensor with values sampled from a standard normal distribution (mean=0, std=1)
using the size entered in the text field.
For example, if you want to create a 4D tensor of batch size=2, channel=3, height=96, width=32, enter:
[2,3,96,32]
category: PyTorch wrapper - Tensor creation
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"size": ("STRING", {"multiline": True}),
"data_type": (("float32", "float16", "bfloat16", "float64"),)
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
@classmethod
def IS_CHANGED(cls, **kw):
return float("NaN")
def f(self, size: str, data_type: str) -> tuple:
"""
Creates a PyTorch tensor with values sampled from a standard normal distribution.
Args:
size (str): String value in the text field.
data_type (str): Desired floating-point data type of the tensor.
Returns:
tuple: A tuple containing the PyTorch tensor.
"""
# Convert string to Python list
list_size = ast.literal_eval(size)
# Convert list to PyTorch Size
sz = torch.Size(list_size)
tens = torch.randn(sz, dtype=DTYPE_MAPPING[data_type])
return (tens,)
modules/pytorch_wrapper/pt_remainder.py
-41
View File
@@ -1,41 +0,0 @@
from typing import Any, Dict
import torch
class PtRemainder:
"""
Computes the element-wise remainder of division between two PyTorch tensors.
category: PyTorch wrapper - Arithmetic operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens_a": ("TENSOR", {}),
"tens_b": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens_a: torch.Tensor, tens_b: torch.Tensor) -> tuple:
"""
Computes the element-wise remainder of division between two PyTorch tensors.
Args:
tens_a (torch.Tensor): A PyTorch tensor of any shape.
tens_b (torch.Tensor): A PyTorch tensor of shape that supports element-wise remainder computation with `tens_a`.
Returns:
tuple: A tuple containing the resultant tensor.
"""
return (torch.remainder(tens_a, tens_b),)
modules/pytorch_wrapper/pt_reshape.py
-62
View File
@@ -1,62 +0,0 @@
from typing import Any, Dict
import torch
import ast
class PtReshape:
"""
Reshapes a PyTorch tensor into a specified shape using `torch.reshape()`.
The total number of elements must remain unchanged.
For example, if a tensor has shape (2, 3, 4), and `new_shape="[6, 4]"`,
the output will have shape (6, 4).
Use -1 to automatically infer a dimension, e.g.:
```
[2, -1]
```
category: PyTorch wrapper - Transform
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"tens": ("TENSOR", {}),
"new_shape": ("STRING", {"multiline": True})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, new_shape: str) -> tuple:
"""
Reshapes the input tensor based on the specified shape.
Args:
tens (torch.Tensor): The input tensor.
new_shape (str): A string representation of the desired shape,
e.g., "[2, -1]" or "(6, 4)".
Returns:
tuple: A tuple containing the reshaped PyTorch tensor.
"""
# Convert string to a tuple of integers
shape_order = tuple(ast.literal_eval(new_shape))
# Ensure the parsed shape is valid
if not isinstance(shape_order, tuple) or not all(isinstance(i, int) for i in shape_order):
raise ValueError(f"Invalid format for new_shape: {new_shape}. Must be a tuple or list of integers.")
# Apply reshape operation
tens = tens.reshape(*shape_order)
return (tens,)
modules/pytorch_wrapper/pt_save_model.py
-36
View File
@@ -1,36 +0,0 @@
import os
from typing import Any, Dict
import torch
from .utils import get_model_full_path
class PtSaveModel:
"""
A wrapper class for saving a PyTorch model.
category: IO
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"model": ("PTMODEL", {}),
"model_path": ("STRING", {"default": "model.pkl"})
}
}
RETURN_NAMES: tuple = ("model full path",)
RETURN_TYPES: tuple = ("STRING",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, model:Any, model_path: str):
model_full_path = get_model_full_path(model_path)
torch.save(model.state_dict(), model_full_path)
return (os.path.realpath(model_full_path),)
modules/pytorch_wrapper/pt_scatter.py
-69
View File
@@ -1,69 +0,0 @@
from typing import Any, Dict
import torch
import ast
class PtScatter:
"""
Generates a new tensor by replacing values at specified positions using an index tensor.
* Target tensor (tens): Some of whose values are replaced.
* Source tensor (src): Values from this tensor replace a part or whole of the target tensor.
* Index tensor (index): Specifies the positions in the target tensor for replacement.
category: PyTorch wrapper - Indexing and Slicing Operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"tens": ("TENSOR", {}),
"dim": ("INT", {"default": 0, "min": -10, "max": 10}),
"index": ("STRING", {"multiline": True}),
"src": ("STRING", {"multiline": True}),
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self,
tens: torch.Tensor,
dim: int,
index: str,
src: str) -> tuple:
"""
Replaces elements in the target tensor along the specified dimension using an index tensor.
Args:
tens (torch.Tensor): The target tensor whose values will be replaced.
dim (int): The dimension along which values are replaced.
index (str): A string representation of the index tensor (list format).
src (str): A string representation of the source tensor (list format).
Returns:
tuple: A tuple containing the modified tensor.
"""
# Parse the input index and source tensors from string representations
index_list = ast.literal_eval(index)
src_list = ast.literal_eval(src)
# Convert to tensors
index_tensor = torch.tensor(index_list, dtype=torch.long, device=tens.device)
src_tensor = torch.tensor(src_list, dtype=tens.dtype, device=tens.device)
# Ensure index and src shapes match
if index_tensor.shape != src_tensor.shape:
raise ValueError(f"Shape mismatch: index tensor {index_tensor.shape} and src tensor {src_tensor.shape} must match.")
# Perform the scatter operation
updated_tens = torch.scatter(tens.clone(), dim=dim, index=index_tensor, src=src_tensor)
return (updated_tens,)
modules/pytorch_wrapper/pt_show_size.py
-60
View File
@@ -1,60 +0,0 @@
"""
Almost all of this code was taken from
https://github.com/pythongosssss/ComfyUI-Custom-Scripts/blob/main/py/show_text.py
See credit/credit.md for the full license.
"""
from typing import List
import torch
class PtShowSize:
"""
Displays PyTorch Size object as a string.
category: Display data
"""
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"size": ("PTSIZE", {}),
},
"hidden": {
"unique_id": "UNIQUE_ID",
"extra_pnginfo": "EXTRA_PNGINFO",
},
}
INPUT_IS_LIST = True
RETURN_TYPES = ("STRING",)
FUNCTION = "notify"
OUTPUT_NODE = True
OUTPUT_IS_LIST = (True,)
CATEGORY = "Data Analysis"
MAX_NUM_ELEMENTS_TO_DISPLAY = 100
def notify(self, size: torch.Size, unique_id=None, extra_pnginfo=None):
text = str(size[0])
if unique_id is not None and extra_pnginfo is not None:
if not isinstance(extra_pnginfo, list):
print("Error: extra_pnginfo is not a list")
elif (
not isinstance(extra_pnginfo[0], dict)
or "workflow" not in extra_pnginfo[0]
):
print("Error: extra_pnginfo[0] is not a dict or missing 'workflow' key")
else:
workflow = extra_pnginfo[0]["workflow"]
node = next(
(x for x in workflow["nodes"] if str(x["id"]) == str(unique_id[0])),
None,
)
if node:
node["widgets_values"] = [text]
else:
print("unique_id or extra_pnginfo is None.")
return {"ui": {"text": [text]}, "result": ([text],)}
modules/pytorch_wrapper/pt_show_text.py
-62
View File
@@ -1,62 +0,0 @@
"""
Almost all of this code was taken from
https://github.com/pythongosssss/ComfyUI-Custom-Scripts/blob/main/py/show_text.py
See credit/credit.md for the full license.
"""
from typing import List
import torch
class PtShowText:
"""
Displays PyTorch tensor as a string. Note that the tensor is partially printed out when
the tensor contains more elements than the threshold set internally.
category: Display data
"""
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"tens": ("TENSOR", {}),
},
"hidden": {
"unique_id": "UNIQUE_ID",
"extra_pnginfo": "EXTRA_PNGINFO",
},
}
INPUT_IS_LIST = True
RETURN_TYPES = ("STRING",)
FUNCTION = "notify"
OUTPUT_NODE = True
OUTPUT_IS_LIST = (True,)
CATEGORY = "Data Analysis"
MAX_NUM_ELEMENTS_TO_DISPLAY = 100
def notify(self, tens: torch.Tensor, unique_id=None, extra_pnginfo=None):
torch.set_printoptions(threshold=PtShowText.MAX_NUM_ELEMENTS_TO_DISPLAY)
text = str(tens[0])
if unique_id is not None and extra_pnginfo is not None:
if not isinstance(extra_pnginfo, list):
print("Error: extra_pnginfo is not a list")
elif (
not isinstance(extra_pnginfo[0], dict)
or "workflow" not in extra_pnginfo[0]
):
print("Error: extra_pnginfo[0] is not a dict or missing 'workflow' key")
else:
workflow = extra_pnginfo[0]["workflow"]
node = next(
(x for x in workflow["nodes"] if str(x["id"]) == str(unique_id[0])),
None,
)
if node:
node["widgets_values"] = [text]
else:
print("unique_id or extra_pnginfo is None.")
return {"ui": {"text": [text]}, "result": ([text],)}
modules/pytorch_wrapper/pt_sin.py
-39
View File
@@ -1,39 +0,0 @@
from typing import Any, Dict
import torch
class PtSin:
"""
Computes the sine of a PyTorch tensor element-wise.
category: PyTorch wrapper - Math operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens_a": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens_a: torch.Tensor) -> tuple:
"""
Computes the element-wise sine of a PyTorch tensor.
Args:
tens_a (torch.Tensor): A PyTorch tensor of numerical type, representing angles in radians.
Returns:
tuple: A tuple containing the resultant tensor after applying the sine function.
"""
return (torch.sin(tens_a),)
modules/pytorch_wrapper/pt_sinh.py
-39
View File
@@ -1,39 +0,0 @@
from typing import Any, Dict
import torch
class PtSinh:
"""
Computes the hyperbolic sine of a PyTorch tensor element-wise.
category: PyTorch wrapper - Math operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens_a": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens_a: torch.Tensor) -> tuple:
"""
Computes the element-wise hyperbolic sine of a PyTorch tensor.
Args:
tens_a (torch.Tensor): A PyTorch tensor of numerical type.
Returns:
tuple: A tuple containing the resultant tensor after applying the hyperbolic sine function.
"""
return (torch.sinh(tens_a),)
modules/pytorch_wrapper/pt_size.py
-41
View File
@@ -1,41 +0,0 @@
from typing import Any, Dict
import torch
import ast
class PtSize:
"""
Extracts the PyTorch Size object of a PyTorch tensor using the size() method.
category: PyTorch wrapper - Size object support
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"tens": ("TENSOR",),
}
}
RETURN_TYPES: tuple = ("PTSIZE",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor) -> tuple:
"""
Extracts the PyTorch Size object of a PyTorch tensor using the size() method.
Args:
tens (torch.Tensor): The tensor.
Returns:
tuple: A tuple containing the PyTorch Size.
"""
return (tens.size(),)
modules/pytorch_wrapper/pt_size_create.py
-46
View File
@@ -1,46 +0,0 @@
from typing import Any, Dict
import torch
import ast
class PtSizeCreate:
"""
Creates a PyTorch Size using values entered in the text field.
category: PyTorch wrapper - Size object support
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"data": ("STRING", {"multiline": True}),
}
}
RETURN_TYPES: tuple = ("PTSIZE",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, data: str) -> tuple:
"""
Creates a PyTorch Size using values entered in the text field.
Args:
data (str): String value in the text field.
Returns:
tuple: A tuple containing the PyTorch Size.
"""
# Convert string to Python list
list_data = ast.literal_eval(data)
# Convert list to PyTorch Size
size = torch.Size(list_data)
return (size,)
modules/pytorch_wrapper/pt_size_to_numpy.py
-41
View File
@@ -1,41 +0,0 @@
from typing import Any, Dict
import torch
import numpy as np
class PtSizeToNumpy:
"""
Converts PyTorch Size object to NumPy ndarray.
category: PyTorch wrapper - Size object support
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"sz": ("PTSIZE", {}),
}
}
RETURN_TYPES: tuple = ("NDARRAY",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, sz: torch.Size) -> tuple:
"""
Converts PyTorch Size object to NumPy ndarray.
Args:
sz (torch.Size): PyTorch Size
Returns:
tuple: A tuple containing the ndarray.
"""
array = np.array(sz, dtype=int)
return (array,)
modules/pytorch_wrapper/pt_size_to_string.py
-40
View File
@@ -1,40 +0,0 @@
from typing import Any, Dict
import torch
class PtSizeToString:
"""
Converts PyTorch Size object to a Python string.
category: PyTorch wrapper - Size object support
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"sz": ("PTSIZE", {}),
}
}
RETURN_TYPES: tuple = ("STRING",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, sz: torch.Size) -> tuple:
"""
Converts PyTorch Size object to a Python string.
Args:
sz (torch.Size): PyTorch Size
Returns:
tuple: A tuple containing the string.
"""
return (str(sz),)
modules/pytorch_wrapper/pt_sqrt.py
-39
View File
@@ -1,39 +0,0 @@
from typing import Any, Dict
import torch
class PtSqrt:
"""
Computes the square root of each element in a PyTorch tensor.
category: PyTorch wrapper - Math operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens_a": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens_a: torch.Tensor) -> tuple:
"""
Computes the square root of a PyTorch tensor element-wise.
Args:
tens_a (torch.Tensor): A PyTorch tensor of numerical type.
Returns:
tuple: A tuple containing the resultant tensor after applying the square root function.
"""
return (torch.sqrt(tens_a),)
modules/pytorch_wrapper/pt_squeeze.py
-45
View File
@@ -1,45 +0,0 @@
from typing import Any, Dict
import torch
from .utils import DTYPE_MAPPING
class PtSqueeze:
"""
Removes a dimension at the specified position in the input tensor if it is of size 1.
Specifying a tuple of dimensions is not supported yet.
category: PyTorch wrapper - Transform
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"tens": ("TENSOR", {}),
"dim": ("INT", {"default": -1, "min": -1, "max": 10})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, dim: int) -> tuple:
"""
Removes a dimension at the specified position in the input tensor if it is of size 1.
Args:
tens (torch.Tensor): Input tensor.
dim (int): Dimension to remove if it is of size 1.
Returns:
tuple: A tuple containing the transformed PyTorch tensor.
"""
tens = torch.squeeze(tens, dim)
return (tens,)
modules/pytorch_wrapper/pt_std.py
-68
View File
@@ -1,68 +0,0 @@
from typing import Any, Dict
import torch
from .utils import str_to_dim
class PtStd:
"""
Computes the standard deviation of a PyTorch tensor along the specified dimension(s).
Specify 1 to use Bessel's correction (N-1) for sample standard deviation.
Specify 0 to compute population standard deviation.
Specify the dimension(s) in the `dim` field using an integer, a list, or a tuple, as shown below:
```
0, [0], or (1, 2)
```
category: PyTorch wrapper - Reduction operation & Summary statistics
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"tens": ("TENSOR", {}),
"correction": ("INT", {"default": 1, "min": 0, "max": 1}),
"dim": ("STRING", {"multiline": True}),
"keepdim": ("BOOLEAN", {"default": False})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, dim: str, correction: int, keepdim: bool) -> tuple:
"""
Computes the standard deviation along the specified dimension(s).
Args:
tens (torch.Tensor): The input tensor.
dim (str): A string representation of the dimension(s) along which to compute the standard deviation,
e.g., "0", "[0]", or "(1, 2)".
correction (int): 0 for population standard deviation, 1 for sample standard deviation (Bessel's correction).
keepdim (bool): Whether to retain reduced dimensions.
Returns:
tuple: A tuple containing the standard deviation tensor.
"""
dim = str_to_dim(dim)
# Ensure correction is valid (0 or 1)
if correction not in (0, 1):
raise ValueError(f"Invalid correction value: {correction}. Must be 0 (population) or 1 (sample).")
# Ensure keepdim is a boolean
if not isinstance(keepdim, bool):
raise TypeError(f"Invalid type for keepdim: {keepdim}. Must be a boolean.")
# Compute the standard deviation along the specified dimension(s)
tens = torch.std(tens, dim=dim, correction=correction, keepdim=keepdim)
return (tens,)
modules/pytorch_wrapper/pt_sub.py
-41
View File
@@ -1,41 +0,0 @@
from typing import Any, Dict
import torch
class PtSub:
"""
Subtracts one PyTorch tensor from another.
category: PyTorch wrapper - Arithmetic operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens_a": ("TENSOR", {}),
"tens_b": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens_a: torch.Tensor, tens_b: torch.Tensor) -> tuple:
"""
Subtracts one PyTorch tensor from another.
Args:
tens_a (torch.Tensor): A PyTorch tensor of any shape.
tens_b (torch.Tensor): A PyTorch tensor of shape that supports subtraction from `tens_a`.
Returns:
tuple: A tuple containing the resultant tensor.
"""
return (torch.sub(tens_a, tens_b),)
modules/pytorch_wrapper/pt_sum.py
-59
View File
@@ -1,59 +0,0 @@
from typing import Any, Dict
import torch
from .utils import str_to_dim
class PtSum:
"""
Computes the sum of a PyTorch tensor along the specified dimension(s).
Specify the dimension(s) in the `dim` field using an integer, a list, or a tuple, as shown below:
```
0, [0], or (1, 2)
```
category: PyTorch wrapper - Reduction operation & Summary statistics
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"tens": ("TENSOR", {}),
"dim": ("STRING", {"multiline": True}),
"keepdim": ("BOOLEAN", {"default": False})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, dim: str, keepdim: bool) -> tuple:
"""
Computes the sum along the specified dimension(s).
Args:
tens (torch.Tensor): The input tensor.
dim (str): A string representation of the dimension(s) along which to compute the sum,
e.g., "0", "[0]", or "(1, 2)".
keepdim (bool): Whether to retain reduced dimensions.
Returns:
tuple: A tuple containing the summed tensor.
"""
dim = str_to_dim(dim)
# Ensure keepdim is a boolean
if not isinstance(keepdim, bool):
raise TypeError(f"Invalid type for keepdim: {keepdim}. Must be a boolean.")
# Compute the sum along the specified dimension(s)
tens = torch.sum(tens, dim=dim, keepdim=keepdim)
return (tens,)
modules/pytorch_wrapper/pt_tan.py
-39
View File
@@ -1,39 +0,0 @@
from typing import Any, Dict
import torch
class PtTan:
"""
Computes the tangent of a PyTorch tensor element-wise.
category: PyTorch wrapper - Math operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens_a": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens_a: torch.Tensor) -> tuple:
"""
Computes the element-wise tangent of a PyTorch tensor.
Args:
tens_a (torch.Tensor): A PyTorch tensor of numerical type, representing angles in radians.
Returns:
tuple: A tuple containing the resultant tensor after applying the tangent function.
"""
return (torch.tan(tens_a),)
modules/pytorch_wrapper/pt_tanh.py
-39
View File
@@ -1,39 +0,0 @@
from typing import Any, Dict
import torch
class PtTanh:
"""
Computes the hyperbolic tangent of a PyTorch tensor element-wise.
category: PyTorch wrapper - Math operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens_a": ("TENSOR", {})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens_a: torch.Tensor) -> tuple:
"""
Computes the element-wise hyperbolic tangent of a PyTorch tensor.
Args:
tens_a (torch.Tensor): A PyTorch tensor of numerical type.
Returns:
tuple: A tuple containing the resultant tensor after applying the hyperbolic tangent function.
"""
return (torch.tanh(tens_a),)
modules/pytorch_wrapper/pt_to_image.py
-41
View File
@@ -1,41 +0,0 @@
from typing import Any, Dict
import torch
class PtToImage:
"""
Casts a PyTorch tensor as an Image tensor.
category: PyTorch wrapper - Image processing
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"tens": ("TENSOR", {}),
}
}
RETURN_NAMES: tuple = ("image",)
RETURN_TYPES: tuple = ("IMAGE",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor) -> tuple:
"""
Casts a PyTorch tensor as an Image tensor.
Args:
tens (torch.Tensor): PyTorch Tensor
Returns:
tuple: A tuple containing the image tensor.
"""
return (tens,)
modules/pytorch_wrapper/pt_to_image_transpose.py
-47
View File
@@ -1,47 +0,0 @@
from typing import Any, Dict
import torch
class PtToImageTranspose:
"""
Casts a PyTorch tensor as an Image tensor and transposes it from (C, H, W) to (H, W, C). For rank-4 inputs, the batch axis remains unchanged.
category: PyTorch wrapper - Image processing
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"tens": ("TENSOR", {}),
}
}
RETURN_NAMES: tuple = ("image",)
RETURN_TYPES: tuple = ("IMAGE",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor) -> tuple:
"""
Casts a PyTorch tensor as an Image tensor.
Args:
tens (torch.Tensor): PyTorch Tensor
Returns:
tuple: A tuple containing the image tensor.
"""
if tens.dim() == 3:
image = tens.permute(1, 2, 0)
elif tens.dim() == 4:
image = tens.permute(0, 2, 3, 1)
else:
raise ValueError("Only rank 3 or 4 tensors are supported.")
return (image,)
modules/pytorch_wrapper/pt_to_latent.py
-40
View File
@@ -1,40 +0,0 @@
from typing import Any, Dict
import torch
class PtToLatent:
"""
Casts a PyTorch tensor as a latent tensor.
category: PyTorch wrapper - Tensor data conversion
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"tens": ("TENSOR", {}),
}
}
RETURN_TYPES: tuple = ("LATENT",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor) -> tuple:
"""
Casts a PyTorch tensor as a latent tensor.
Args:
tens (torch.Tensor): PyTorch Tensor
Returns:
tuple: A tuple containing the latent tensor.
"""
return ({"samples":tens},)
modules/pytorch_wrapper/pt_to_numpy.py
-41
View File
@@ -1,41 +0,0 @@
from typing import Any, Dict
import torch
class PtToNumpy:
"""
Converts PyTorch tensor to NumPy ndarray.
category: PyTorch wrapper - Tensor data conversion
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"tens": ("TENSOR", {}),
}
}
RETURN_TYPES: tuple = ("NDARRAY",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor) -> tuple:
"""
Converts PyTorch tensor to NumPy ndarray.
Args:
tens (torch.Tensor): PyTorch Tensor
Returns:
tuple: A tuple containing the ndarray.
"""
array = tens.cpu().detach().numpy()
return (array,)
modules/pytorch_wrapper/pt_to_rgb_tensors.py
-50
View File
@@ -1,50 +0,0 @@
from typing import Any, Dict
import torch
class PtToRgbTensors:
"""
Splits a PyTorch tensor into R, G, and B tensors.
category: PyTorch wrapper - Tensor data conversion
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the required input types.
Returns:
Dict[str, Any]: A dictionary of required input types.
"""
return {
"required": {
"tens": ("TENSOR", {}),
}
}
RETURN_NAMES: tuple = ("R", "G", "B")
RETURN_TYPES: tuple = ("TENSOR", "TENSOR", "TENSOR")
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor) -> tuple:
"""
Splits a PyTorch tensor into R, G, and B tensors.
Args:
tens (torch.Tensor): A PyTorch tensor of shape [B, H, W, C] or [H, W, C].
Assumes the last axis represents the color channels.
Returns:
tuple: A tuple containing the R, G, and B tensors.
If the input tensor has a batch axis, it will be retained.
Raises:
ValueError: If the input tensor does not have exactly 3 channels.
"""
if tens.size()[-1] != 3:
raise ValueError(f"Expected a 3-channel tensor, but got {tens.size()[-1]} channels.")
r, g, b = tens[..., 0], tens[..., 1], tens[..., 2]
return (r, g, b)
modules/pytorch_wrapper/pt_train_classification_model.py
-122
View File
@@ -1,122 +0,0 @@
from typing import Any, Dict, Optional, Tuple
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
class PtTrainClassificationModel:
"""
Trains a classification model using a given dataset, optimizer, and number of epochs.
category: PyTorch wrapper - Training
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types required for training the model.
Returns:
Dict[str, Any]: A dictionary specifying required and optional input types.
"""
return {
"required": {
"model": ("PTMODEL", {}),
"train_loader": ("PTDATALOADER", {}),
"optimizer": ("PTOPTIMIZER", {}),
"epochs": ("INT", {"default":1, "min":1, "max":1e6}),
"use_gpu": ("BOOLEAN", {"default":False})
},
"optional": {
"val_loader": ("PTDATALOADER", {}),
}
}
RETURN_NAMES: tuple = ("Model", "train loss", "val loss")
RETURN_TYPES: tuple = ("PTMODEL","TENSOR","TENSOR")
FUNCTION: str = "f"
CATEGORY: str = "Training"
@classmethod
def IS_CHANGED(cls, **kw):
return float("NaN")
def f(self, model: nn.Module,
train_loader: torch.utils.data.DataLoader,
optimizer: torch.optim.Optimizer,
epochs: int,
use_gpu: bool,
val_loader: Optional[torch.utils.data.DataLoader] = None) -> Tuple[nn.Module, torch.Tensor, torch.Tensor]:
"""
Trains a classification model.
Args:
model (torch.nn.Module): The PyTorch model to be trained.
train_loader (torch.utils.data.DataLoader): DataLoader containing the training dataset.
optimizer (torch.optim.Optimizer): Optimizer used for training.
epochs (int): Number of training epochs.
use_gpu (bool): Whether to use GPU for training.
val_loader (Optional[torch.utils.data.DataLoader]): DataLoader for validation dataset (optional).
Returns:
Tuple[torch.nn.Module, torch.Tensor, torch.Tensor]: A tuple containing the trained model,
training loss history, and validation loss history.
"""
with torch.inference_mode(False):
train_loss = []
val_loss = []
if use_gpu:
model.to("cuda")
loss_function = nn.CrossEntropyLoss() # mean loss
# Train the model
for epoch in range(epochs):
model.train()
total_loss = 0
total_samples = len(train_loader.dataset)
for images, labels in train_loader:
if use_gpu:
images = images.to("cuda")
labels = labels.to("cuda")
optimizer.zero_grad()
outputs = model(images)
loss = loss_function(outputs, labels)
loss.backward()
optimizer.step()
batch_size = labels.size(0)
batch_loss = loss.item() * batch_size
total_loss += batch_loss
epoch_loss = total_loss / total_samples
print(f"Epoch (train) {epoch+1}/{epochs}, Loss: {epoch_loss:.4f}")
train_loss.append(epoch_loss)
if val_loader:
with torch.no_grad():
model.train(False) # TODO: Change to model.*e*v*a*l() once Comfy's security checker is fixed.
total_loss = 0
total_samples = len(val_loader.dataset)
for images, labels in val_loader:
if use_gpu:
images = images.to("cuda")
labels = labels.to("cuda")
outputs = model(images)
loss = loss_function(outputs, labels)
batch_size = labels.size(0)
batch_loss = loss.item() * batch_size
total_loss += batch_loss
epoch_loss = total_loss / total_samples
print(f"Epoch (val) {epoch+1}/{epochs}, Loss: {epoch_loss:.4f}")
val_loss.append(epoch_loss)
if use_gpu:
model.to("cpu")
return (model,
torch.tensor(train_loss, dtype=torch.float32),
torch.tensor(val_loss, dtype=torch.float32))
modules/pytorch_wrapper/pt_unsqueeze.py
-44
View File
@@ -1,44 +0,0 @@
from typing import Any, Dict
import torch
class PtUnsqueeze:
"""
Adds a singleton dimension at the specified position in the input tensor.
For example, if a tensor has shape (3, 4) and `dim=1`, the output will have shape (3, 1, 4).
category: PyTorch wrapper - Transform
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"tens": ("TENSOR", {}),
"dim": ("INT", {"default": -1, "min": -10, "max": 10})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, dim: int) -> tuple:
"""
Adds a singleton dimension at the specified position.
Args:
tens (torch.Tensor): Input tensor.
dim (int): Dimension at which to insert the new singleton dimension.
Returns:
tuple: A tuple containing the transformed PyTorch tensor.
"""
tens = torch.unsqueeze(tens, dim)
return (tens,)
modules/pytorch_wrapper/pt_var.py
-68
View File
@@ -1,68 +0,0 @@
from typing import Any, Dict
import torch
from .utils import str_to_dim
class PtVar:
"""
Computes the variance of a PyTorch tensor along the specified dimension(s).
Specify 1 to use Bessel's correction (N-1) for sample variance.
Specify 0 to compute population variance.
Specify the dimension(s) in the `dim` field using an integer, a list, or a tuple, as shown below:
```
0, [0], or (1, 2)
```
category: PyTorch wrapper - Reduction operation & Summary statistics
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"tens": ("TENSOR", {}),
"correction": ("INT", {"default": 1, "min": 0, "max": 1}),
"dim": ("STRING", {"multiline": True}),
"keepdim": ("BOOLEAN", {"default": False})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, dim: str, correction: int, keepdim: bool) -> tuple:
"""
Computes the variance along the specified dimension(s).
Args:
tens (torch.Tensor): The input tensor.
dim (str): A string representation of the dimension(s) along which to compute the variance,
e.g., "0", "[0]", or "(1, 2)".
correction (int): 0 for population variance, 1 for sample variance (Bessel's correction).
keepdim (bool): Whether to retain reduced dimensions.
Returns:
tuple: A tuple containing the variance tensor.
"""
dim = str_to_dim(dim)
# Ensure correction is valid (0 or 1)
if correction not in (0, 1):
raise ValueError(f"Invalid correction value: {correction}. Must be 0 (population) or 1 (sample).")
# Ensure keepdim is a boolean
if not isinstance(keepdim, bool):
raise TypeError(f"Invalid type for keepdim: {keepdim}. Must be a boolean.")
# Compute the variance along the specified dimension(s)
tens = torch.var(tens, dim=dim, correction=correction, keepdim=keepdim)
return (tens,)
modules/pytorch_wrapper/pt_view.py
-62
View File
@@ -1,62 +0,0 @@
from typing import Any, Dict
import torch
import ast
class PtView:
"""
Reshapes a PyTorch tensor into a specified shape using `torch.view()`.
The total number of elements must remain unchanged.
For example, if a tensor has shape (2, 3, 4), and `new_shape="[6, 4]"`,
the output will have shape (6, 4).
Use -1 to automatically infer a dimension, e.g.:
```
[2, -1]
```
category: PyTorch wrapper - Transform
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"tens": ("TENSOR", {}),
"new_shape": ("STRING", {"multiline": True})
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, tens: torch.Tensor, new_shape: str) -> tuple:
"""
Reshapes the input tensor based on the specified shape.
Args:
tens (torch.Tensor): The input tensor.
new_shape (str): A string representation of the desired shape,
e.g., "[2, -1]" or "(6, 4)".
Returns:
tuple: A tuple containing the reshaped PyTorch tensor.
"""
# Convert string to a tuple of integers
shape_order = tuple(ast.literal_eval(new_shape))
# Ensure the parsed shape is valid
if not isinstance(shape_order, tuple) or not all(isinstance(i, int) for i in shape_order):
raise ValueError(f"Invalid format for new_shape: {new_shape}. Must be a tuple or list of integers.")
# Apply view operation
tens = tens.view(*shape_order)
return (tens,)
modules/pytorch_wrapper/pt_where.py
-61
View File
@@ -1,61 +0,0 @@
from typing import Any, Dict
import torch
class PtWhere:
"""
Generates a new tensor by selecting values based on a condition tensor.
* `condition_tens` (torch.Tensor): Boolean tensor that determines which values are selected.
* `true_tens` (torch.Tensor): Values to use where the condition is `True`.
* `false_tens` (torch.Tensor): Values to use where the condition is `False`.
category: PyTorch wrapper - Indexing and Slicing Operations
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"condition_tens": ("TENSOR", {}),
"true_tens": ("TENSOR", {}),
"false_tens": ("TENSOR", {}),
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self,
condition_tens: torch.Tensor,
true_tens: torch.Tensor,
false_tens: torch.Tensor) -> tuple:
"""
Selects values from `true_tens` or `false_tens` based on `condition_tens`.
Args:
condition_tens (torch.Tensor): Boolean tensor indicating selection conditions.
true_tens (torch.Tensor): Values used where condition is True.
false_tens (torch.Tensor): Values used where condition is False.
Returns:
tuple: A tuple containing the modified tensor.
"""
# Ensure condition_tens is a boolean tensor
if condition_tens.dtype != torch.bool:
raise TypeError(f"Expected `condition_tens` to be a boolean tensor, got {condition_tens.dtype}")
# Ensure all tensors are compatible in shape
if not torch.broadcast_shapes(condition_tens.shape, true_tens.shape, false_tens.shape):
raise ValueError(f"Shape mismatch: condition {condition_tens.shape}, true {true_tens.shape}, false {false_tens.shape}")
# Perform the conditional selection
tens_out = torch.where(condition_tens, true_tens, false_tens)
return (tens_out,)
modules/pytorch_wrapper/pt_zeros.py
-52
View File
@@ -1,52 +0,0 @@
from typing import Any, Dict
import torch
import ast
from .utils import DTYPE_MAPPING
class PtZeros:
"""
Creates a PyTorch tensor of zeros using the size entered in the text field.
For example, if you want to create a 4D tensor of batch size=2, channel=3, height=96, width=32, enter:
[2,3,96,32]
category: PyTorch wrapper - Tensor creation
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"size": ("STRING", {"multiline": True}),
"data_type": (("float32", "float16", "bfloat16", "float64", "uint8", "int8", "int16", "int32", "int64", "bool"),)
}
}
RETURN_TYPES: tuple = ("TENSOR",)
FUNCTION: str = "f"
CATEGORY: str = "Data Analysis"
def f(self, size: str, data_type: str) -> tuple:
"""
Creates a PyTorch tensor of zeros using the size entered in the text field.
Args:
size (str): String value in the text field.
data_type (str): Desired data type of the tensor.
Returns:
tuple: A tuple containing the PyTorch tensor.
"""
# Convert string to Python list
list_size = ast.literal_eval(size)
# Convert list to PyTorch Size
sz = torch.Size(list_size)
tens = torch.zeros(sz, dtype=DTYPE_MAPPING[data_type])
return (tens,)
modules/pytorch_wrapper/ptn_conv_model.py
-99
View File
@@ -1,99 +0,0 @@
from typing import Any, Dict, Tuple
import ast
import torch
from .ptn_conv_model_def import ConvModel
class PtnConvModel:
"""
A convolutional model consisting of multiple convolutional layers.
### Fields:
- `input_dim`: A string representing the input dimensions in the format "(C,H,W)".
Example: `"(3,28,28)"` for a 3-channel 28x28 image.
- `penultimate_dim`: An integer specifying the number of features before the output layer. If you specify 0, the number is computed internally.
Default: `0`
- `output_dim`: An integer specifying the number of output features, which should match the number of target classes in classification.
Default: `10`, Min: `1`, Max: `1e6`.
- `channel_list`: A string representing a Python list of integers specifying the number of channels per layer excluding the channel for the input (e.g. 3 for the color image).
Example: `"[32,64,128,256,512]"`.
- `kernel_size_list`: A string representing a Python list of kernel sizes per layer.
Example: `"[3,3,3,3,1]"`.
- `padding_list`: A string representing a Python list of padding values per layer.
Example: `"[1,1,1,1,0]"`.
- `downsample_list`: A string representing a Python list of boolean values indicating whether to downsample at each layer.
Example: `"[True,True,True,True,False]"`.
category: PyTorch wrapper - Training
"""
@classmethod
def INPUT_TYPES(cls) -> Dict[str, Any]:
"""
Defines the input types for the function.
Returns:
Dict[str, Any]: A dictionary specifying required input types.
"""
return {
"required": {
"input_dim": ("STRING", {"default": "(3,28,28)"}),
"penultimate_dim": ("INT", {"default": 0, "min": 0, "max": 1e6}),
"output_dim": ("INT", {"default": 10, "min": 1, "max": 1e6}),
"channel_list": ("STRING", {"default": "[32,64,128,256,512]"}),
"kernel_size_list": ("STRING", {"default": "[3,3,3,3,1]"}),
"padding_list": ("STRING", {"default": "[1,1,1,1,0]"}),
"downsample_list": ("STRING", {"default": "[True,True,True,True,False]"})
}
}
RETURN_TYPES: Tuple[str] = ("PTMODEL",)
FUNCTION: str = "f"
CATEGORY: str = "Training"
@classmethod
def IS_CHANGED(cls, **kw):
return float("NaN")
def f(self,
input_dim:str,
penultimate_dim: int,
output_dim: int,
channel_list: str,
kernel_size_list: str,
padding_list: str,
downsample_list: str) -> None:
"""
Constructs a convolutional neural network model.
Args:
TODO:Update
Returns:
Tuple[nn.Module]: A tuple containing the instantiated PyTorch model.
Raises:
TODO:Update
"""
input_dim_parsed = tuple(ast.literal_eval(input_dim))
channel_list_parsed = ast.literal_eval(channel_list)
kernel_size_list_parsed = ast.literal_eval(kernel_size_list)
padding_list_parsed = ast.literal_eval(padding_list)
downsample_list_parsed = ast.literal_eval(downsample_list)
with torch.inference_mode(False):
model = ConvModel(
input_dim_parsed,
penultimate_dim,
output_dim,
channel_list_parsed,
kernel_size_list_parsed,
padding_list_parsed,
downsample_list_parsed)
return (model,)
modules/pytorch_wrapper/ptn_conv_model_def.py
-151
View File
@@ -1,151 +0,0 @@
from typing import List, Dict, Any, Tuple
import torch
import torch.nn as nn
class ConvSingleLayer(nn.Module):
"""A single convolutional layer followed by an optional ReLU activation and optional downsampling.
Shape is tensor of shape `(batch_size, in_channels, height, width)`.
Args:
in_channels (int): Number of input channels.
out_channels (int): Number of output channels.
kernel_size (int): Size of the convolutional kernel.
padding (int): Padding size for convolution.
downsample (bool): If True, applies a 2x2 MaxPooling layer.
apply_activation (bool, optional): If True, applies ReLU activation. Defaults to True.
pragma: skip_doc
"""
@classmethod # TODO: Dummy method to suppress ComfyUI warning. DO NOT REMOVE
def INPUT_TYPES(cls) -> Dict[str, Any]:
return {}
def __init__(self,
in_channels: int,
out_channels: int,
kernel_size: int,
padding: int,
downsample: bool,
apply_activation: bool = True) -> None:
super().__init__()
self.conv = nn.Conv2d(in_channels=in_channels,
out_channels=out_channels,
kernel_size=kernel_size,
padding=padding)
self.apply_activation = apply_activation
self.activation = nn.ReLU()
self.downsample = downsample
self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2)
def forward(self, inputs: torch.Tensor) -> torch.Tensor:
"""Forward pass of the convolutional layer.
Args:
inputs (torch.Tensor): Input tensor of shape `(batch_size, in_channels, height, width)`.
Returns:
torch.Tensor: Output tensor after convolution, optional activation, and optional pooling.
"""
x = self.conv(inputs)
if self.downsample:
x = self.maxpool(x)
if self.apply_activation:
x = self.activation(x)
return x
class ConvModel(nn.Module):
"""A multi-layer convolutional neural network with ReLU activation after each convolutional layer, except for the final dense layer.
Args:
input_dim (Tuple[int, int, int]): Shape of a single input sample (C, H, W).
penultimate_dim (int): Number of features before the final dense layer. If 0 is specified, this value is automatically computed internally.
output_dim (int): Number of output features.
channel_list (List[int]): List of channel sizes for each layer. This should not include the inputs channel.
Input channel is taken from the first element of input_dim.
kernel_size_list (List[int]): List of kernel sizes for each convolutional layer.
padding_list (List[int]): List of padding values for each convolutional layer.
downsample_list (List[bool]): List indicating whether downsampling should be applied after each layer.
pragma: skip_doc
"""
@classmethod # TODO: Dummy method to suppress ComfyUI warning
def INPUT_TYPES(cls) -> Dict[str, Any]:
return {}
def __init__(self,
input_dim: Tuple[int, int, int],
penultimate_dim: int,
output_dim: int,
channel_list: List[int],
kernel_size_list: List[int],
padding_list: List[int],
downsample_list: List[bool]) -> None:
super().__init__()
# Validate list lengths
if (len(channel_list) != len(kernel_size_list) or
len(kernel_size_list) != len(padding_list) or
len(padding_list) != len(downsample_list)):
raise ValueError("channel_list, kernel_size_list, padding_list, and downsample_list must have the same length.")
self.channel_list = [input_dim[0]] + channel_list
# Initialize convolutional layers
layers = []
for i in range(len(self.channel_list) - 1):
layers.append(
ConvSingleLayer(
in_channels=self.channel_list[i],
out_channels=self.channel_list[i + 1],
kernel_size=kernel_size_list[i],
padding=padding_list[i],
downsample=downsample_list[i]
)
)
self.conv_layers = nn.ModuleList(layers)
if penultimate_dim == 0:
# Try to auto compute penultimate_dim by processing the dummy input
with torch.no_grad():
dummy_input = torch.randn(1, *input_dim) # B, C, H, W
x = dummy_input
for layer in self.conv_layers:
x = layer(x)
checked_penultimate_dim = x.numel() # Number of elements
else:
checked_penultimate_dim = penultimate_dim
self.final_dense = nn.Linear(checked_penultimate_dim, output_dim, bias=True)
def forward(self, inputs: torch.Tensor) -> torch.Tensor:
"""Forward pass of the model.
Args:
inputs (torch.Tensor): Input tensor of shape `(batch_size, C, H, W)`.
Returns:
torch.Tensor: Output tensor of shape `(batch_size, output_dim)`.
Raises:
ValueError: If inputs is not a rank 4 tensor.
"""
if inputs.dim() not in [3, 4]:
raise ValueError("Inputs is not a rank 3 or 4 tensor.")
if inputs.dim() == 3:
x = torch.unsqueeze(inputs, 1)
else:
x = inputs
for layer in self.conv_layers:
x = layer(x)
x = torch.flatten(x, start_dim=1) # Flatten except batch-axis
x = self.final_dense(x)
return x

Some files were not shown because too many files have changed in this diff Show More