Added Ptn Pre Add Channel Axis node and Ptn Pre Flatten node. Also added example workflows to build a model from scratch using Ptn Chained Model node.

__init__.py

@@ -128,6 +128,8 @@ from .modules.pytorch_wrapper.ptn_conv_model import PtnConvModel
 from .modules.pytorch_wrapper.ptn_linear import PtnLinear
 from .modules.pytorch_wrapper.ptn_linear_model import PtnLinearModel
 from .modules.pytorch_wrapper.ptn_max_pool2d import PtnMaxPool2d
+from .modules.pytorch_wrapper.ptn_pre_add_channel_axis import PtnPreAddChannelAxis
+from .modules.pytorch_wrapper.ptn_pre_flatten import PtnPreFlatten
 from .modules.pytorch_wrapper.ptn_resnet_model import PtnResnetModel
 from .modules.pytorch_wrapper.pto_adam import PtoAdam
 from .modules.pytorch_wrapper.pto_adamw import PtoAdamW
@@ -280,6 +282,8 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = {
     "PtnLinear": PtnLinear,
     "PtnLinearModel": PtnLinearModel,
     "PtnMaxPool2d": PtnMaxPool2d,
+    "PtnPreAddChannelAxis": PtnPreAddChannelAxis,
+    "PtnPreFlatten": PtnPreFlatten,
     "PtnResnetModel": PtnResnetModel,
     "PtoAdam": PtoAdam,
     "PtoAdamW": PtoAdamW,
@@ -432,6 +436,8 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = {
     "PtnLinear": "Ptn Linear",
     "PtnLinearModel": "Ptn Linear Model",
     "PtnMaxPool2d": "Ptn Max Pool2d",
+    "PtnPreAddChannelAxis": "Ptn Pre Add Channel Axis",
+    "PtnPreFlatten": "Ptn Pre Flatten",
     "PtnResnetModel": "Ptn Resnet Model",
     "PtoAdam": "Pto Adam",
     "PtoAdamW": "Pto AdamW",

docs/building_a_model_from_scratch.md

@@ -0,0 +1,12 @@
+# Building a Model from Scratch
+
+ComfyUI Pt Wrapper provides composite models that can be easily customized by adjusting parameters. However, you can also build a model from scratch by sequentially combining basic models such as `Conv2D` and `Linear`.
+
+To chain two models, use the **Pt Chained Model** node. This node also supports specifying post-processing after the forward pass of the second model (e.g., activation functions). To define the post-processing step, connect a **Ptf** node (e.g., `Ptf ReLU`).
+
+In some cases, preprocessing may be required before feeding data into the model. For instance, you can use the **Pt Pre Add Channel Axis** node to add a channel dimension to the input tensor or the **Pt Pre Flatten** node to flatten convolutional output before passing it to a linear layer.
+
+The following example workflows demonstrate this approach:
+
+- `./examples/workflows/fashion_mnist_eval_conv_chained_model.json`: Constructs a multi-layer convolutional network from scratch.
+- `./examples/workflows/fashion_mnist_train_chained_model.json`: Uses the trained model for evaluation.

docs/model_training.md

@@ -7,13 +7,12 @@ Currently, following archituctures are supported (node names are in parentheses)
 * Convolutional Network (Ptn Conv Model)
 * ResNet (Ptn Resnet Model)
 
-In addition, a single layer Linear model is supported as an example to illustrate the use (see the bottom of this page on how to use it):
-* Linear (Ptn Linear)
-
 You can adjust the model configuration including output dimension on the UI.
 
 The system is designed to be modular. For example, to switch from a linear model to a convolutional model, you only need to replace the model node (see screenshot below).
 
+You can also build a model from scratch by sequentially combining basic models such as `Conv2D` and `Linear` by using the **Pt Chained Model** node.  See [Building a Model from Scratch](building_a_model_from_scratch.md) for details.
+
 You can choose to:
 * Use your own images
 * Use a public dataset

docs/reference/node_reference.md

@@ -120,6 +120,8 @@
 | [Ptn Linear](ptn_linear.md) | A linear model consisting of a single dense layer.   |
 | [Ptn Linear Model](ptn_linear_model.md) | A linear model consisting of dense layers.   |
 | [Ptn Max Pool](ptn_max_pool2d.md) | A maxpool layer.   |
+| [Ptn Pre Add Channel Axis](ptn_pre_add_channel_axis.md) | Adds a channel axis after the batch axis if the input is rank 3 (bs, h, w) |
+| [Ptn Pre Flatten](ptn_pre_flatten.md) | Flattens the input tensor before processing the tensor in the specified model. |
 | [Ptn Resnet Model](ptn_resnet_model.md) | A Resnet model consisting of multiple Resnet layers.   |
 ## PyTorch wrapper - Optimizer
 | Node | Description |

docs/reference/ptn_conv2d.md

@@ -25,7 +25,7 @@ t.
 | dilation | String |
 | groups | Int |
 | bias | Boolean |
-| padding_mode | Zeros |
+| padding_mode |  |
 
 ## Output
 | Data type |

docs/reference/ptn_pre_add_channel_axis.md

@@ -0,0 +1,18 @@
+# Ptn Pre Add Channel Axis
+Adds a channel axis after the batch axis if the input is rank 3 (bs, h, w)
+before the forward pass by the specified model.
+
+## Input
+| Name | Data type |
+|---|---|
+| model | Ptmodel |
+
+## Output
+| Data type |
+|---|
+| Ptmodel |
+
+<HR>
+Category: PyTorch wrapper - Model
+
+ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.

docs/reference/ptn_pre_flatten.md

@@ -0,0 +1,17 @@
+# Ptn Pre Flatten
+Flattens the input tensor before processing the tensor in the specified model.
+
+## Input
+| Name | Data type |
+|---|---|
+| model | Ptmodel |
+
+## Output
+| Data type |
+|---|
+| Ptmodel |
+
+<HR>
+Category: PyTorch wrapper - Model
+
+ComfyUI Pt Wrapper Node Reference. © 2025 Hide Inada (HowToSD.com). All rights reserved.

modules/pytorch_wrapper/ptn_conv_model_def.py

@@ -29,12 +29,13 @@ class ConvSingleLayer(nn.Module):
                  kernel_size: int,
                  padding: int,
                  downsample: bool,
-                 bias:bool = False,
+                 bias:bool = True,
                  apply_activation: bool = True) -> None:
         super().__init__()
         self.conv = nn.Conv2d(in_channels=in_channels,
                               out_channels=out_channels,
                               kernel_size=kernel_size,
+                              bias=bias,
                               padding=padding)
         self.apply_activation = apply_activation
         self.activation = nn.ReLU()

modules/pytorch_wrapper/ptn_pre_add_channel_axis.py

@@ -0,0 +1,84 @@
+from typing import Any, Dict, Tuple, Callable, Optional
+import torch
+import torch.nn as nn
+
+
+class PtnPreAddChannelAxisDef(nn.Module):
+    """
+    Adds a channel axis after the batch axis if the input is rank 3 (bs, h, w).
+
+    pragma: skip_doc
+    """
+
+    def __init__(self,
+                 model: nn.Module
+    ):
+        super().__init__()
+        self.model = model
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass of the model.
+
+        Args:
+            inputs (torch.Tensor): Input tensor.
+
+        Returns:
+            torch.Tensor: Output tensor of the `model`.
+        """
+        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
+
+        x = self.model(x)
+        return x
+
+
+class PtnPreAddChannelAxis:
+    """
+    Adds a channel axis after the batch axis if the input is rank 3 (bs, h, w)
+    before the forward pass by the specified model.
+
+    category: PyTorch wrapper - Model
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
+        """
+        Defines the expected input types.
+
+        Returns:
+            Dict[str, Any]: A dictionary specifying required and optional input types.
+        """
+        return {
+            "required": {
+                "model": ("PTMODEL", {})
+            }
+        }
+
+    RETURN_TYPES: Tuple[str, ...] = ("PTMODEL",)
+    FUNCTION: str = "f"
+    CATEGORY: str = "Training"
+
+    @classmethod
+    def IS_CHANGED(cls, **kw):
+        return float("NaN")
+
+    def f(self,
+          model: nn.Module):
+        """
+        Adds a channel axis after the batch axis if the input is rank 3 (bs, h, w).
+
+        Args:
+            model (nn.Module): The model.
+
+        Returns:
+            Tuple[nn.Module]: A tuple containing the instantiated model.
+        """
+        with torch.inference_mode(False):
+            output_model = PtnPreAddChannelAxisDef(model)
+        return (output_model,)

modules/pytorch_wrapper/ptn_pre_flatten.py

@@ -0,0 +1,76 @@
+from typing import Any, Dict, Tuple
+import torch
+import torch.nn as nn
+
+
+class PtnPreFlattenDef(nn.Module):
+    """
+    Flattens the input tensor before processing the tensor in the specified model.
+
+    pragma: skip_doc
+    """
+
+    def __init__(self,
+                 model: nn.Module
+    ):
+        super().__init__()
+        self.model = model
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass of the model.
+
+        Args:
+            inputs (torch.Tensor): Input tensor.
+
+        Returns:
+            torch.Tensor: Output tensor of the `model`.
+        """
+        x = torch.flatten(inputs, start_dim=1)  # Flatten except batch-axis
+        x = self.model(x)
+        return x
+
+
+class PtnPreFlatten:
+    """
+    Flattens the input tensor before processing the tensor in the specified model.
+
+    category: PyTorch wrapper - Model
+    """
+
+    @classmethod
+    def INPUT_TYPES(cls) -> Dict[str, Any]:
+        """
+        Defines the expected input types.
+
+        Returns:
+            Dict[str, Any]: A dictionary specifying required and optional input types.
+        """
+        return {
+            "required": {
+                "model": ("PTMODEL", {})
+            }
+        }
+
+    RETURN_TYPES: Tuple[str, ...] = ("PTMODEL",)
+    FUNCTION: str = "f"
+    CATEGORY: str = "Training"
+
+    @classmethod
+    def IS_CHANGED(cls, **kw):
+        return float("NaN")
+
+    def f(self,
+          model: nn.Module):
+        """
+        Adds a channel axis after the batch axis if the input is rank 3 (bs, h, w).
+
+        Args:
+            model (nn.Module): The model.
+
+        Returns:
+            Tuple[nn.Module]: A tuple containing the instantiated model.
+        """
+        with torch.inference_mode(False):
+            output_model = PtnPreFlattenDef(model)
+        return (output_model,)

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "pt-wrapper"
 description = "PyTorch extension for ComfyUI featuring extensive PyTorch wrapper nodes for seamless tensor operations and PyTorch model training."
-version = "3.0.12"
+version = "3.0.13"
 license = {file = "LICENSE"}
 
 [project.urls]

tests/pytorch_wrapper/ptn_pre_add_channel_axis_test.py

@@ -0,0 +1,49 @@
+import os
+import sys
+import unittest
+import torch
+import torch.nn as nn
+
+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 pytorch_wrapper.ptn_pre_add_channel_axis import PtnPreAddChannelAxis
+
+class DummyModel(nn.Module):
+    def forward(self, x):
+        return x  # Pass-through model for testing
+
+
+class TestPtnPreAddChannelAxis(unittest.TestCase):
+    def setUp(self):
+        """Set up test instance."""
+        self.model_node = PtnPreAddChannelAxis()
+        self.model = self.model_node.f(model=DummyModel())[0]
+        
+
+    def test_valid_shapes(self):
+        """Test valid input shapes (rank 3 and 4)."""
+        rank3_tensor = torch.randn(2, 8, 8)  # (bs, h, w)
+        rank4_tensor = torch.randn(2, 3, 8, 8)  # (bs, c, h, w)
+
+        result3 = self.model(rank3_tensor)
+        result4 = self.model(rank4_tensor)
+
+        self.assertEqual(result3.shape, (2, 1, 8, 8), f"Unexpected shape for rank 3: {result3.shape}")
+        self.assertEqual(result4.shape, (2, 3, 8, 8), f"Unexpected shape for rank 4: {result4.shape}")
+
+    def test_invalid_shapes(self):
+        """Test invalid input shapes (rank 2 and rank 5)."""
+        rank2_tensor = torch.randn(8, 8)  # (h, w)
+        rank5_tensor = torch.randn(2, 3, 8, 8, 8)  # (bs, c, d, h, w)
+
+        with self.assertRaises(ValueError, msg="Rank 2 tensor did not raise ValueError"):
+            self.model(rank2_tensor)
+
+        with self.assertRaises(ValueError, msg="Rank 5 tensor did not raise ValueError"):
+            self.model(rank5_tensor)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/pytorch_wrapper/ptn_pre_flatten_test.py

@@ -0,0 +1,40 @@
+import os
+import sys
+import unittest
+import torch
+import torch.nn as nn
+
+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 pytorch_wrapper.ptn_pre_flatten import PtnPreFlatten
+
+class DummyModel(nn.Module):
+    def forward(self, x):
+        return x  # Pass-through model for testing
+
+
+class TestPtnPreFlatten(unittest.TestCase):
+    def setUp(self):
+        """Set up test instance."""
+        self.model_node = PtnPreFlatten()
+        self.model = self.model_node.f(model=DummyModel())[0]
+
+    def test_valid_shapes(self):
+        """Test flattening on valid input shapes."""
+        rank3_tensor = torch.randn(2, 8, 8)  # (bs, h, w)
+        rank4_tensor = torch.randn(2, 3, 8, 8)  # (bs, c, h, w)
+        rank5_tensor = torch.randn(2, 3, 4, 8, 8)  # (bs, c, d, h, w)
+
+        result3 = self.model(rank3_tensor)
+        result4 = self.model(rank4_tensor)
+        result5 = self.model(rank5_tensor)
+
+        self.assertEqual(result3.shape, (2, 64), f"Unexpected shape for rank 3: {result3.shape}")
+        self.assertEqual(result4.shape, (2, 192), f"Unexpected shape for rank 4: {result4.shape}")
+        self.assertEqual(result5.shape, (2, 768), f"Unexpected shape for rank 5: {result5.shape}")
+
+
+if __name__ == "__main__":
+    unittest.main()