Added more Pt nodes.

modules/pytorch_wrapper/pt_flatten.py

@@ -0,0 +1,40 @@
+from typing import Any, Dict
+import torch
+
+
+class PtFlatten:
+    """
+    Flattens a PyTorch tensor into a 1D tensor.
+
+    Category: PyTorch wrapper
+    """
+    
+    @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 = "PyTorch wrapper"
+
+    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_from_latent.py

@@ -37,4 +37,4 @@ class PtFromLatent:
         Returns:
             tuple: A tuple containing the PyTorch tensor.
         """
-        return (latent,)
+        return (latent["samples"],)

modules/pytorch_wrapper/pt_size.py

@@ -0,0 +1,41 @@
+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
+    """
+    
+    @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 = "PyTorch wrapper"
+
+    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

@@ -0,0 +1,46 @@
+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
+    """
+    
+    @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 = "PyTorch wrapper"
+
+    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

@@ -0,0 +1,41 @@
+from typing import Any, Dict
+import torch
+import numpy as np
+
+class PtSizeToNumpy:
+    """
+    Converts PyTorch Size object to NumPy ndarray.
+
+    category: PyTorch wrapper
+    """
+    
+    @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 = "PyTorch wrapper"
+
+    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

@@ -0,0 +1,40 @@
+from typing import Any, Dict
+import torch
+
+
+class PtSizeToString:
+    """
+    Converts PyTorch Size object to a Python string.
+
+    category: PyTorch wrapper
+    """
+    
+    @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 = "PyTorch wrapper"
+
+    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_to_latent.py

@@ -37,4 +37,4 @@ class PtToLatent:
         Returns:
             tuple: A tuple containing the latent tensor.
         """
-        return (tens,)
+        return ({"samples":tens},)

modules/pytorch_wrapper/pt_to_rgb_tensors.py

@@ -0,0 +1,50 @@
+from typing import Any, Dict
+import torch
+
+
+class PtToRgbTensors:
+    """
+    Splits a PyTorch tensor into R, G, and B tensors.
+
+    Category: PyTorch wrapper
+    """
+    
+    @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 = "PyTorch wrapper"
+
+    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)