diff --git a/__init__.py b/__init__.py index 3010ba7..c38a468 100644 --- a/__init__.py +++ b/__init__.py @@ -101,11 +101,6 @@ from .modules.pytorch_wrapper.pt_zeros import PtZeros from .modules.pytorch_wrapper.ptn_conv_model import PtnConvModel from .modules.pytorch_wrapper.ptn_linear_model import PtnLinearModel from .modules.pytorch_wrapper.ptn_resnet_model import PtnResnetModel -from .modules.pytorch_wrapper.ptn_resnet_model_def import NResnetBlocks -from .modules.pytorch_wrapper.ptn_resnet_model_def import ResnetModel -from .modules.pytorch_wrapper.ptn_resnet_model_def import SingleConvBlock -from .modules.pytorch_wrapper.ptn_resnet_model_def import SingleConvBlockNoActivation -from .modules.pytorch_wrapper.ptn_resnet_model_def import SingleResnetBlock from .modules.pytorch_wrapper.pto_adam import PtoAdam from .modules.pytorch_wrapper.ptv_dataset import PtvDataset from .modules.pytorch_wrapper.ptv_dataset_len import PtvDatasetLen @@ -122,7 +117,6 @@ NODE_CLASS_MAPPINGS (Dict[str, Type[T]]): """ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = { - "NResnetBlocks": NResnetBlocks, "PtAbs": PtAbs, "PtAcos": PtAcos, "PtAdd": PtAdd, @@ -231,10 +225,6 @@ NODE_CLASS_MAPPINGS: Dict[str, Type[T]] = { "PtvImageFolderDataset": PtvImageFolderDataset, "PtvTransformsResize": PtvTransformsResize, "PtvTransformsToTensor": PtvTransformsToTensor, - "ResnetModel": ResnetModel, - "SingleConvBlock": SingleConvBlock, - "SingleConvBlockNoActivation": SingleConvBlockNoActivation, - "SingleResnetBlock": SingleResnetBlock, } @@ -244,7 +234,6 @@ NODE_DISPLAY_NAME_MAPPINGS (Dict[str, str]): """ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = { - "NResnetBlocks": "N Resnet Blocks", "PtAbs": "Pt Abs", "PtAcos": "Pt Acos", "PtAdd": "Pt Add", @@ -353,10 +342,6 @@ NODE_DISPLAY_NAME_MAPPINGS: Dict[str, str] = { "PtvImageFolderDataset": "Ptv Image Folder Dataset", "PtvTransformsResize": "Ptv Transforms Resize", "PtvTransformsToTensor": "Ptv Transforms To Tensor", - "ResnetModel": "Resnet Model", - "SingleConvBlock": "Single Conv Block", - "SingleConvBlockNoActivation": "Single Conv Block No Activation", - "SingleResnetBlock": "Single Resnet Block", } diff --git a/docs/reference/node_reference.md b/docs/reference/node_reference.md index 54898f6..0027797 100644 --- a/docs/reference/node_reference.md +++ b/docs/reference/node_reference.md @@ -8,14 +8,6 @@ | Node | Description | | --- | --- | | [Pt Save Model](pt_save_model.md) | A wrapper class for saving a PyTorch model. | -## Miscellaneous -| Node | Description | -| --- | --- | -| [N Resnet Blocks](ptn_resnet_model_def.md) | | -| [Resnet Model](ptn_resnet_model_def.md) | | -| [Single Conv Block](ptn_resnet_model_def.md) | | -| [Single Conv Block No Activation](ptn_resnet_model_def.md) | | -| [Single Resnet Block](ptn_resnet_model_def.md) | | ## PyTorch wrapper - Arithmetic operations | Node | Description | | --- | --- | diff --git a/modules/pytorch_wrapper/ptn_resnet_model_def.py b/modules/pytorch_wrapper/ptn_resnet_model_def.py index 4bb56ca..dc61e6b 100644 --- a/modules/pytorch_wrapper/ptn_resnet_model_def.py +++ b/modules/pytorch_wrapper/ptn_resnet_model_def.py @@ -8,6 +8,10 @@ import torch.nn as nn import torch.nn.functional as F class SingleConvBlock(nn.Module): + """ + pragma: skip_doc + """ + def __init__(self, in_channels, out_channels, kernel_size, strides): super().__init__() padding = kernel_size // 2 @@ -33,6 +37,10 @@ class SingleConvBlock(nn.Module): class SingleConvBlockNoActivation(nn.Module): + """ + pragma: skip_doc + """ + def __init__(self, in_channels, out_channels, kernel_size, strides): super().__init__() padding = kernel_size // 2 @@ -54,6 +62,9 @@ class SingleConvBlockNoActivation(nn.Module): class SingleResnetBlock(nn.Module): + """ + pragma: skip_doc + """ def __init__(self, in_channels, out_channels, kernel_size=3, downsample=False): super().__init__() @@ -120,7 +131,10 @@ class SingleResnetBlock(nn.Module): class NResnetBlocks(nn.Module): - + """ + pragma: skip_doc + """ + def __init__( self, num_blocks, @@ -168,6 +182,10 @@ class NResnetBlocks(nn.Module): return tens class ResnetModel(nn.Module): + """ + pragma: skip_doc + """ + def __init__(self, num_blocks, in_channels, diff --git a/pyproject.toml b/pyproject.toml index 06564c2..3d6e1eb 100644 --- a/pyproject.toml +++ b/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.3" +version = "3.0.4" license = {file = "LICENSE"} [project.urls]