TashaSkyUp-EternalKernelPytorchNodes/pytorch_nodes.py

import torch
# import folder_paths
import torchvision
import torch.nn as nn
import functools
import torch
import json

try:
    from .config import config_settings
except ImportError as e:
    from config import config_settings

from . import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS


def ETK_pytorch_base(cls):
    cls.CATEGORY = "ETK/pytorch"
    # Add spaces to the camel case class name
    pretty_name = cls.__name__
    for i in range(1, len(pretty_name)):
        if pretty_name[i].isupper():
            pretty_name = pretty_name[:i] + " " + pretty_name[i:]
    cls.DISPLAY_NAME = pretty_name
    NODE_CLASS_MAPPINGS[cls.__name__] = cls
    NODE_DISPLAY_NAME_MAPPINGS[cls.DISPLAY_NAME] = pretty_name

    # Wrap the function defined in the FUNCTION attribute
    func_name = getattr(cls, "FUNCTION", None)
    if func_name and hasattr(cls, func_name):
        original_func = getattr(cls, func_name)

        @functools.wraps(original_func)
        def wrapped_func(*args, **kwargs):
            with torch.inference_mode(False):
                return original_func(*args, **kwargs)

        setattr(cls, func_name, wrapped_func)

    return cls


@ETK_pytorch_base
class PyTorchDatasetDownloader:
    # Define the input types for the node
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "dataset_name": ("STRING", {"default": "mnist"}),
                "download_path": ("STRING", {"default": "./data"}),
                "mode": (["dataset", "tensors", ],),
            }
        }

    # Define the return types
    RETURN_TYPES = ("TORCH_DATASET", "TORCH_TENSOR", "TORCH_TENSOR",)
    RETURN_NAMES = ("Full Dataset", "Features tensor", "labels tensor",)
    FUNCTION = "download_dataset"

    # Method to download the dataset
    def download_dataset(self, dataset_name, download_path, mode):
        from torchvision import datasets, transforms
        import torch.utils.data as data_utils

        # check if the dataset already exists at the download_path

        # Downloading the dataset
        if dataset_name.lower() == "mnist":
            exists = False
            try:
                dataset = datasets.MNIST(
                    root=download_path,
                    train=True,
                    transform=transforms.ToTensor(),
                    download=False
                )
                exists = True
            except:
                dataset = datasets.MNIST(
                    root=download_path,
                    train=True,
                    transform=transforms.ToTensor(),
                    download=True
                )
        else:
            raise ValueError("Unsupported dataset")

        if mode == "tensors":
            # Split dataset into inputs and labels tensors
            inputs_tensor, labels_tensor = self.split_dataset_if_labels_exist(dataset)

            # Return tensors
            return (None, inputs_tensor, labels_tensor,)
        else:
            return (dataset, None, None,)

    def split_dataset_if_labels_exist(self, dataset):
        try:
            item = dataset[0]
            if isinstance(item, tuple) and len(item) > 1:
                # Extract inputs and labels
                inputs = [data[0] for data in dataset]
                labels = [data[1] for data in dataset]

                return torch.stack(inputs), torch.tensor(labels)
            else:
                return dataset, None
        except (TypeError, IndexError):
            return dataset, None


@ETK_pytorch_base
class DatsetSplitter:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "dataset": ("TORCH_DATASET",),
            }
        }

    RETURN_TYPES = ("TORCH_DATASET", "TORCH_DATASET")
    RETURN_NAMES = ("train_ds", "val_ds")
    FUNCTION = "split_dataset"

    def split_dataset(self, dataset):
        import torch
        from torch.utils.data import DataLoader
        from torch.utils.data import random_split

        ds_len = len(dataset)
        val_ratio = 0.2
        # use random_split to create a validation set
        train_ds, val_ds = random_split(dataset, [int(ds_len * (1 - val_ratio)), int(ds_len * val_ratio)])

        # train_ds is now of type torch.utils.data.Dataset throw an error if it is not
        if not isinstance(train_ds, torch.utils.data.Dataset):
            raise TypeError("train_ds is not of type torch.utils.data.Dataset")

        # val_ds is now of type torch.utils.data.Dataset throw an error if it is not
        if not isinstance(val_ds, torch.utils.data.Dataset):
            raise TypeError("val_ds is not of type torch.utils.data.Dataset")

        return (train_ds, val_ds,)


@ETK_pytorch_base
class AddLinearLayerNode:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model": ("TORCH_MODEL",),
                "in_features": ("INT", {"min": 1, "max": 2 ** 24}),
                "out_features": ("INT", {"min": 1, "max": 1_000_000_000}),
                "bias": ([True, False],),
                "initialization": (["default", "xavier_uniform", "xavier_normal"],),
                "dtype": (["float32", "float64", "float16", "int32", "int64", "int16", "int8", "uint8"],),
            }
        }

    RETURN_TYPES = ("TORCH_MODEL",)
    FUNCTION = "add_linear_layer"
    CATEGORY = "model"

    def add_linear_layer(self, model, in_features, out_features, bias=True, initialization="default", dtype="float32"):
        import torch
        # Create the new linear layer
        use_dtype = getattr(torch, dtype)

        new_linear_layer = nn.Linear(in_features, out_features, bias=bias, dtype=use_dtype)
        new_linear_layer.requires_grad_(True)

        # Initialize the weights if required
        if initialization == "xavier_uniform":
            nn.init.xavier_uniform_(new_linear_layer.weight)
        elif initialization == "xavier_normal":
            nn.init.xavier_normal_(new_linear_layer.weight)

        # Reconstruct the nn.Sequential model to include the new layer
        if isinstance(model, nn.Sequential):
            # layers = list(model.children())
            # layers.append(new_linear_layer)
            # new_model = nn.Sequential(*layers)
            model.insert(len(model), new_linear_layer)
        else:
            raise TypeError("The provided model is not a nn.Sequential model.")

        # Returning the modified model
        return (model,)


@ETK_pytorch_base
class AddSoftmaxLayerNode:
    """
    SOFTMAX is a function that turns a vector of K real values into a vector of K real values that sum to 1.
    """

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model": ("TORCH_MODEL",),
                "dim": ("INT", {"min": 0}),
            },
            "optional": {
                "log_softmax": ([False, True],),
            }

        }

    RETURN_TYPES = ("TORCH_MODEL",)
    FUNCTION = "add_softmax_layer"
    CATEGORY = "model"

    def add_softmax_layer(self, model, dim, log_softmax=False):
        import torch
        # Create the new softmax layer
        if log_softmax:
            new_softmax_layer = nn.LogSoftmax(dim=dim)
        else:
            new_softmax_layer = nn.Softmax(dim=dim)

        # Reconstruct the nn.Sequential model to include the new layer
        if isinstance(model, nn.Sequential):
            # layers = list(model.children())
            # layers.append(new_linear_layer)
            # new_model = nn.Sequential(*layers)
            model.insert(len(model), new_softmax_layer)
        else:
            raise TypeError("The provided model is not a nn.Sequential model.")

        # Returning the modified model
        return (model,)


@ETK_pytorch_base
class AddConvLayer:
    """
    Convolutional layer
    """

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model": ("TORCH_MODEL",),
                "one_d": ("BOOLEAN", {"default": False}),
                "in_channels": ("INT", {"min": 1}),
                "out_channels": ("INT", {"min": 1}),
                "kernel_size": ("INT", {"min": 1}),
                "stride": ("INT", {"min": 1}),
                "padding": ("STRING", {"default":"(0,0)"}),
                "bias": ([True, False],),
                "initialization": (["default", "xavier_uniform", "xavier_normal"],),
                "dtype": (["float32", "float64", "float16", "int32", "int64", "int16", "int8", "uint8"],),
            }
        }

    RETURN_TYPES = ("TORCH_MODEL",)
    FUNCTION = "add_conv_layer"
    CATEGORY = "model"

    def add_conv_layer(self, model, one_d, in_channels, out_channels, kernel_size, stride, padding, bias=True,
                       initialization="default", dtype="float32"):
        import torch
        # Create the new convolutional layer
        use_dtype = getattr(torch, dtype)
        padding=padding.replace("(", "")
        padding=padding.replace(")", "")
        padding = tuple(map(int, padding.split(",")))
        if not one_d:
            new_conv_layer = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=bias,
                                       dtype=use_dtype)
        else:
            new_conv_layer = nn.Conv1d(in_channels, out_channels, kernel_size, stride, padding, bias=bias,
                                       dtype=use_dtype)

        new_conv_layer.requires_grad_(True)

        # Initialize the weights if required
        if initialization == "xavier_uniform":
            nn.init.xavier_uniform_(new_conv_layer.weight)
        elif initialization == "xavier_normal":
            nn.init.xavier_normal_(new_conv_layer.weight)

        # Reconstruct the nn.Sequential model to include the new layer
        if isinstance(model, nn.Sequential):
            # layers = list(model.children())
            # layers.append(new_linear_layer)
            # new_model = nn.Sequential(*layers)
            model.insert(len(model), new_conv_layer)
        else:
            raise TypeError("The provided model is not a nn.Sequential model.")

        # Returning the modified model
        return (model,)


@ETK_pytorch_base
class SequentialModelProvider:
    # Define the input types for the node
    @classmethod
    def INPUT_TYPES(cls):
        # Inputs can be used to customize the default architecture
        return {
            "required": {
                "name": ("STRING", {"default": "sequential_model"}),
            }
        }

    FUNCTION = "provide_default_model"
    # Define the return types
    RETURN_TYPES = ("TORCH_MODEL",)

    # Method to provide the default sequential model
    def provide_default_model(self, name):
        # Logic to create the default sequential model based on input configuration
        model = nn.Sequential()

        return (model,)


@ETK_pytorch_base
class PyTorchInferenceNode:
    """
    allows users to perform inference using a pre-trained PyTorch model
    on a given tensor or dataset
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": ("TORCH_MODEL",),  # The pre-trained PyTorch model
            },
            "optional": {
                "input_data": ("TORCH_TENSOR",),  # The input data to perform inference on
                "dataset": ("TORCH_DATASET",),  # The dataset to perform inference on
            }
        }

    RETURN_TYPES = ("TORCH_TENSOR",)
    FUNCTION = "inference"

    def inference(self, model, input_data=None, dataset=None):
        # perform inference on the input data or dataset using the model
        # and return the output
        if input_data != None and dataset != None:
            raise ValueError("Only one of input_data or dataset can be provided")
        elif input_data is not None:
            x = input_data
        elif dataset is not None:
            x = dataset.dataset
            # prepare x so that it can be passed to the model
            # e.g. convert to torch.Tensor
            # avoid the error: ValueError: only one element tensors can be converted to Python scalars
            if isinstance(x, torch.utils.data.Dataset):
                x = next(iter(x))[0]
            elif isinstance(x, torch.utils.data.DataLoader):
                x = next(iter(x))[0]
            elif isinstance(x, torch.Tensor):
                pass



        else:
            raise ValueError("Either input_data or dataset must be provided")

        mdl_device = next(model.parameters()).device
        # check that x and the model are on the same device
        if x.device != mdl_device:
            x = x.to(mdl_device)
        output = model(x)
        return (output,)


@ETK_pytorch_base
class FlattenTensor:
    """
    allows users to flatten a tensor
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "input_data": ("TORCH_TENSOR",),  # The input data to perform inference on
            }
        }

    RETURN_TYPES = ("TORCH_TENSOR",)
    FUNCTION = "flatten"

    def flatten(self, input_data):
        # perform inference on the input data or dataset using the model
        # and return the output
        output = input_data.flatten()
        return (output,)


@ETK_pytorch_base
class FlattenDataset:
    """
    allows users to flatten a dataset
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "dataset": ("TORCH_DATASET",),  # The input data to perform inference on
            }
        }

    RETURN_TYPES = ("TORCH_TENSOR",)
    FUNCTION = "flatten"

    def flatten(self, dataset):
        # perform inference on the input data or dataset using the model
        # and return the output
        output = torch.flatten(dataset.data)
        return (output,)


@ETK_pytorch_base
class ReshapeTensor:
    """
    allows users to reshape a tensor
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "input_data": ("TORCH_TENSOR",),  # The input data to perform inference on
                "shape": ("STRING", {"default": "1, -1"}),
            }
        }

    RETURN_TYPES = ("TORCH_TENSOR",)
    FUNCTION = "reshape"

    def reshape(self, input_data, shape):
        # perform inference on the input data or dataset using the model
        # and return the output
        shape = tuple(map(int, shape.split(",")))
        output = input_data.reshape(shape)
        return (output,)


@ETK_pytorch_base
class ChangeTensorType:
    """
    allows users to change the type of a tensor
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "input_data": ("TORCH_TENSOR",),  # The input data to perform inference on
                "dtype": ("STRING", {"default": "torch.float32"}),
            }
        }

    RETURN_TYPES = ("TORCH_TENSOR",)
    FUNCTION = "change_type"

    def change_type(self, input_data, dtype):
        # perform inference on the input data or dataset using the model
        # and return the output
        dtype_map = {
            "torch.float32": torch.float32,
            "torch.float64": torch.float64,
            "torch.int32": torch.int32,
            "torch.int64": torch.int64,
            "torch.bool": torch.bool,
            "torch.uint8": torch.uint8,
            "torch.int8": torch.int8,
            "torch.int16": torch.int16,
            "torch.half": torch.half,
            "torch.bfloat16": torch.bfloat16,
        }
        if dtype in dtype_map:
            dtype = dtype_map[dtype]
        else:
            # Fallback for any torch.dtype attributes
            if dtype.startswith("torch."):
                dtype_name = dtype.split(".")[-1]
                dtype = getattr(torch, dtype_name, torch.float32)
            else:
                dtype = torch.float32
        output = input_data.type(dtype)
        return (output,)


@ETK_pytorch_base
class TensorToList:
    """
    allows users to convert a tensor to a list
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "input_data": ("TORCH_TENSOR",),  # The input data to perform inference on
            }
        }

    RETURN_TYPES = ("LIST",)
    FUNCTION = "tolist"

    def tolist(self, input_data):
        # perform inference on the input data or dataset using the model
        # and return the output
        output = input_data.tolist()
        return (output,)


@ETK_pytorch_base
class SliceTensor:
    """
    allows users to slice a tensor
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "input_data": ("TORCH_TENSOR",),  # The input data to perform inference on
                "start": ("STRING", {"default": "0"}),
                "end": ("STRING", {"default": "5"}),
                "step": ("STRING", {"default": "1"}),
            }
        }

    RETURN_TYPES = ("TORCH_TENSOR",)
    FUNCTION = "slice"

    def slice(self, input_data, start, end, step):
        # perform inference on the input data or dataset using the model
        # and return the output
        start = int(start)
        end = int(end)
        step = int(step)
        output = input_data[start:end:step]
        return (output,)


@ETK_pytorch_base
class SaveModel:
    """
    allows users to save a model
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": ("TORCH_MODEL",),  # The input data to perform inference on
                "path": ("STRING", {"default": ".data/models/model.pt"}),
            }
        }

    RETURN_TYPES = ("STRING",)
    FUNCTION = "save"
    # define that this is an output
    OUTPUT_NODE = True

    def save(self, model, path):
        import os
        # perform inference on the input data or dataset using the model
        # and return the output
        # create the directory if it doesn't exist
        os.makedirs(os.path.dirname(path), exist_ok=True)
        torch.save(model, path)
        return (path,)

    @classmethod
    def IS_CHANGED(s, **kwargs):
        import hashlib
        data = str(kwargs)
        m = hashlib.sha256()
        m.update(data.encode("utf-8"))
        return m.digest().hex()


@ETK_pytorch_base
class LoadModel:
    """
    allows users to load a model
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "path": ("STRING", {"default": ".data/models/model.pt"}),
            }
        }

    RETURN_TYPES = ("TORCH_MODEL",)
    FUNCTION = "load"

    def load(self, path):
        # perform inference on the input data or dataset using the model
        # and return the output
        print("loaded model: ", path)
        model = torch.load(path)
        return (model,)

    @classmethod
    def IS_CHANGED(s, path):
        # get the date modified of the model at the folder_path
        import os
        import hashlib

        data = str(os.path.getmtime(path))
        data = f"{path} {data}"
        m = hashlib.sha256()
        m.update(data.encode("utf-8"))
        return m.digest().hex()


def optimizer_to(optim, device):
    for param in optim.state.values():
        # Not sure there are any global tensors in the state dict
        if isinstance(param, torch.Tensor):
            param.data = param.data.to(device)
            if param._grad is not None:
                param._grad.data = param._grad.data.to(device)
        elif isinstance(param, dict):
            for subparam in param.values():
                if isinstance(subparam, torch.Tensor):
                    subparam.data = subparam.data.to(device)
                    if subparam._grad is not None:
                        subparam._grad.data = subparam._grad.data.to(device)


@ETK_pytorch_base
class TrainModel:
    """
    allows users to train a model
    """

    @classmethod
    def INPUT_TYPES(cls):
        # Gather available loss functions
        loss_functions = [attr for attr in dir(torch.nn) if attr.endswith("Loss") and callable(getattr(torch.nn, attr))]
        # Gather available optimizers
        optimizers = [attr for attr in dir(torch.optim) if attr[0].isupper() and callable(getattr(torch.optim, attr))]
        # Gather available metrics (for now, just use loss functions + 'accuracy')
        metrics = ["accuracy"] + loss_functions
        return {
            "required": {
                "model": ("TORCH_MODEL",),

            },
            "optional": {
                "dataset": ("TORCH_DATASET",),
                "features tensor": ("TORCH_TENSOR",),
                "labels tensor": ("TORCH_TENSOR",),
                "epochs": ("INT", {"default": 1, "min": 1, "max": 2 ** 24}),
                "batch_size": ("INT", {"default": 1}),
                "loss_function": (loss_functions, {"default": "CrossEntropyLoss"}),
                "optimizer": (optimizers, {"default": "SGD"}),
                "metrics": (metrics, {"default": "accuracy"}),
                "device": ("STRING", {"default": "torch.device('cpu')"}),
                "create_samples": (["TRUE", "FALSE"], {"default": "FALSE"}),
            }
        }

    RETURN_TYPES = ("TORCH_MODEL", "LIST", "TORCH_MODEL")
    RETURN_NAMES = ("model", "metrics", "best model")
    FUNCTION = "train"

    def reset_optimizer_weights(self, optimizer):
        # Loop over the optimizer state
        for param_group in optimizer.param_groups:
            for param in param_group['params']:
                # Reset state dictionary for each parameter
                optimizer.state[param] = torch.optim.Optimizer.StateDict()

    def train(self, model, **kwargs):
        # Consistent boolean handling for create_samples
        create_samples = kwargs.get("create_samples", "FALSE")
        if isinstance(create_samples, str):
            create_samples = create_samples.upper() == "TRUE"
        else:
            create_samples = bool(create_samples)

        with torch.inference_mode(False):
            train_dataset = kwargs.get("dataset", None)
            features_tensor = kwargs.get("features tensor", None)
            labels_tensor = kwargs.get("labels tensor", None)
            if train_dataset is None:
                if features_tensor is None or labels_tensor is None:
                    raise ValueError("Must provide either a dataset or features tensor and labels tensor")
                else:
                    train_dataset = torch.utils.data.TensorDataset(features_tensor, labels_tensor)

            # Consistent int conversion for epochs and batch_size
            epochs = kwargs.get("epochs", 1)
            batch_size = kwargs.get("batch_size", 1)
            try:
                epochs = int(epochs)
            except Exception:
                epochs = 1
            try:
                batch_size = int(batch_size)
            except Exception:
                batch_size = 1

            # Consistent string handling for loss_function, optimizer, metrics
            loss_function = kwargs.get("loss_function", "CrossEntropyLoss")
            optimizer_str = kwargs.get("optimizer", "SGD")
            metrics_str = kwargs.get("metrics", "CrossEntropyLoss")

            # Evaluate loss function
            if isinstance(loss_function, str):
                loss_map = {
                    "MSELoss": torch.nn.MSELoss(),
                    "CrossEntropyLoss": torch.nn.CrossEntropyLoss(),
                    "BCELoss": torch.nn.BCELoss(),
                    "L1Loss": torch.nn.L1Loss(),
                    "NLLLoss": torch.nn.NLLLoss(),
                    "SmoothL1Loss": torch.nn.SmoothL1Loss(),
                    "KLDivLoss": torch.nn.KLDivLoss(),
                    "HuberLoss": torch.nn.HuberLoss(),
                    "torch.nn.MSELoss()": torch.nn.MSELoss(),
                    "torch.nn.CrossEntropyLoss()": torch.nn.CrossEntropyLoss(),
                    "torch.nn.BCELoss()": torch.nn.BCELoss(),
                    "torch.nn.L1Loss()": torch.nn.L1Loss(),
                    "torch.nn.NLLLoss()": torch.nn.NLLLoss(),
                    "torch.nn.SmoothL1Loss()": torch.nn.SmoothL1Loss(),
                    "torch.nn.KLDivLoss()": torch.nn.KLDivLoss(),
                    "torch.nn.HuberLoss()": torch.nn.HuberLoss(),
                }
                if loss_function in loss_map:
                    loss_function = loss_map[loss_function]
                else:
                    # Fallback - try to get from torch.nn
                    clean_name = loss_function.replace("torch.nn.", "").replace("()", "")
                    if hasattr(torch.nn, clean_name):
                        loss_function = getattr(torch.nn, clean_name)()
                    else:
                        loss_function = torch.nn.MSELoss()  # Safe default

            # Evaluate optimizer
            import inspect
            optimizer = None
            if isinstance(optimizer_str, str):
                # Parse optimizer with parameters safely
                if "(" in optimizer_str and ")" in optimizer_str:
                    # Extract optimizer name and parameters
                    opt_name = optimizer_str.split("(")[0].replace("torch.optim.", "")
                    # Use default parameters for safety
                    opt_defaults = {
                        "Adam": {"lr": 0.001, "betas": (0.9, 0.999), "eps": 1e-08},
                        "SGD": {"lr": 0.01, "momentum": 0.9},
                        "RMSprop": {"lr": 0.01, "alpha": 0.99},
                        "AdamW": {"lr": 0.001, "betas": (0.9, 0.999), "eps": 1e-08},
                        "Adagrad": {"lr": 0.01},
                        "Adadelta": {"lr": 1.0, "rho": 0.9},
                    }
                    if opt_name in opt_defaults and hasattr(torch.optim, opt_name):
                        opt_class = getattr(torch.optim, opt_name)
                        optimizer = opt_class(model.parameters(), **opt_defaults[opt_name])
                    else:
                        optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
                else:
                    # Simple name, e.g. 'Adam' or 'SGD'
                    opt_class = getattr(torch.optim, optimizer_str, torch.optim.Adam)
                    sig = inspect.signature(opt_class.__init__)
                    opt_kwargs = {"lr": 0.001}
                    if "momentum" in sig.parameters:
                        opt_kwargs["momentum"] = 0.9
                    optimizer = opt_class(model.parameters(), **opt_kwargs)
            else:
                optimizer = optimizer_str  # fallback

            # Evaluate metrics
            metrics = None
            if isinstance(metrics_str, str):
                if metrics_str.lower() == "accuracy":
                    def accuracy_fn(outputs, labels):
                        # For classification: outputs are logits or probabilities
                        if outputs.dim() > 1 and outputs.size(1) > 1:
                            preds = outputs.argmax(dim=1)
                        else:
                            preds = (outputs > 0.5).long()
                        if labels.dim() > 1 and labels.size(1) > 1:
                            labels = labels.argmax(dim=1)
                        return (preds == labels).float().mean().item()
                    metrics = accuracy_fn
                elif not metrics_str.startswith("torch.nn."):
                    # Safe mapping for common metrics
                    metrics_map = {
                        "L1Loss": torch.nn.L1Loss(),
                        "MSELoss": torch.nn.MSELoss(),
                        "CrossEntropyLoss": torch.nn.CrossEntropyLoss(),
                        "BCELoss": torch.nn.BCELoss(),
                    }
                    if metrics_str in metrics_map:
                        metrics = metrics_map[metrics_str]
                    elif hasattr(torch.nn, metrics_str):
                        metrics = getattr(torch.nn, metrics_str)()
                    else:
                        metrics = torch.nn.MSELoss()  # Safe default
                else:
                    clean_name = metrics_str.replace("torch.nn.", "").replace("()", "")
                    if hasattr(torch.nn, clean_name):
                        metrics = getattr(torch.nn, clean_name)()
                    else:
                        metrics = torch.nn.MSELoss()  # Safe default
            else:
                metrics = metrics_str  # fallback

            train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=0)

            for param in model.parameters():
                param.requires_grad = True

            metrics_out = []
            device = kwargs.get("device", "torch.device('cpu')")
            if isinstance(device, str):
                device_map = {
                    "torch.device('cpu')": torch.device('cpu'),
                    "torch.device('cuda')": torch.device('cuda'),
                    "cpu": torch.device('cpu'),
                    "cuda": torch.device('cuda'),
                    "cuda:0": torch.device('cuda:0'),
                    "cuda:1": torch.device('cuda:1'),
                }
                e_device = device_map.get(device, torch.device('cpu'))
            else:
                e_device = device

            optimizer_to(optimizer, e_device)
            model.to(e_device)
            torch.set_grad_enabled(True)
            torch.cuda.empty_cache()
            best_l = 1024 ** 3
            best_epoch = 0

            import uuid
            import os
            tmpdir = config_settings.get("tmp_dir", "./temp")
            try:
                os.makedirs(tmpdir, exist_ok=True)
            except PermissionError:
                tmpdir = "./temp"
                os.makedirs(tmpdir, exist_ok=True)
            tmp = os.path.join(tmpdir, str(uuid.uuid4()))

            for epoch in range(epochs):
                tot_loss = 0
                for batch_features, batch_labels in train_dataloader:
                    batch_features = batch_features.to(e_device)
                    batch_labels = batch_labels.to(e_device)
                    optimizer.zero_grad()
                    outputs = model(batch_features)
                    loss = loss_function(outputs, batch_labels)
                    tot_loss += loss.item()
                    loss.backward()
                    optimizer.step()
                l = metrics(outputs, batch_labels)
                metrics_out.append(l)
                if os.path.exists(tmp):
                    os.remove(tmp)
                torch.save(model, tmp, _use_new_zipfile_serialization=True, pickle_protocol=4)

                if tot_loss < best_l:
                    if create_samples:
                        import numpy as np
                        import PIL
                        from PIL import Image
                        best_l = tot_loss
                        out_X = None
                        for i in range(len(outputs)):
                            out_tensor = self.make_image_tensor(outputs, i)
                            if out_X is None:
                                out_X = out_tensor
                            else:
                                out_X = np.concatenate((out_X, out_tensor), axis=1)
                        out_y = None
                        for i in range(len(outputs)):
                            label_Tensor = self.make_image_tensor(batch_labels, i)
                            if out_y is None:
                                out_y = label_Tensor
                            else:
                                out_y = np.concatenate((out_y, label_Tensor), axis=1)
                        out_tensor = np.concatenate((out_X, out_y), axis=0)
                        mse = np.mean((out_X - out_y) ** 2)
                        out_tensor = PIL.Image.fromarray(out_tensor)
                        out_tensor.save(f"out{i}_{mse}_{str(loss.item())}.png")
                if os.path.exists(tmp):
                    best_model = torch.load(tmp, weights_only=False)
            # Clean return, no debug prints
            return (model, metrics_out, best_model,)


# --- BEGIN RESTORED NODES FROM 5486f39 ---

@ETK_pytorch_base
class TensorsToDataset:
    """
    allows users to convert a tensor to a dataset
    """
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "input_data": ("TORCH_TENSOR",),
                "labels": ("TORCH_TENSOR",),
            },
            "optional": {
                "move to device": (["cuda", "cpu"], {"default": "cpu"}),
            }
        }
    RETURN_TYPES = ("TORCH_DATASET",)
    FUNCTION = "to_dataset"
    def to_dataset(self, input_data, labels, **kwargs):
        if input_data.is_leaf:
            input_data.requires_grad = True
        try:
            if labels.is_leaf:
                labels.requires_grad = True
        except RuntimeError as e:
            if "floating point and complex dtype" not in str(e):
                raise e
        move_to_device = kwargs.get("move to device", "cpu")
        device_to_use = torch.device(move_to_device)
        input_data = input_data.to(device_to_use)
        labels = labels.to(device_to_use)
        output = torch.utils.data.TensorDataset(input_data, labels)
        return (output,)

@ETK_pytorch_base
class DatasetToDataloader:
    """
    allows users to convert a dataset to a dataloader
    """
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "dataset": ("TORCH_DATASET",),
            }
        }
    RETURN_TYPES = ("TORCH_DATALOADER",)
    FUNCTION = "to_dataloader"
    def to_dataloader(self, dataset):
        output = torch.utils.data.DataLoader(dataset)
        return (output,)

@ETK_pytorch_base
class ComfyUIImageToPytorchTENSOR:
    """
    this just renames the object for compatibility with comfyui->pytorch
    """
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
            }
        }
    RETURN_TYPES = ("TORCH_TENSOR", "TORCH_TENSOR",)
    RETURN_NAMES = ("b,h,w,c", "b,h,w,mean(c)",)
    FUNCTION = "to_tensor"
    def to_tensor(self, image):
        out1 = torch.tensor(image)
        out2 = torch.tensor(image).mean(dim=3)
        return (out1, out2,)

@ETK_pytorch_base
class ListToTensor:
    """
    this just renames the object for compatibility with comfyui->pytorch
    """
    @classmethod
    def INPUT_TYPES(cls):
        valid_dtypes_str = [str(dt) for k, dt in torch.__dict__.items() if isinstance(dt, torch.dtype)]
        return {
            "required": {
                "list": ("LIST",),
            },
            "optional": {
                "dtype": (valid_dtypes_str,),
            }
        }
    RETURN_TYPES = ("TORCH_TENSOR",)
    FUNCTION = "to_tensor"
    def to_tensor(self, list, dtype=str(torch.float32)):
        valid_dtypes_str = [str(dt) for k, dt in torch.__dict__.items() if isinstance(dt, torch.dtype)]
        valid_dtypes = [dt for k, dt in torch.__dict__.items() if isinstance(dt, torch.dtype)]
        if dtype not in valid_dtypes_str:
            raise ValueError(f"Unsupported dtype: {dtype}")
        dto = valid_dtypes[valid_dtypes_str.index(dtype)]
        output = torch.tensor(list, dtype=dto)
        return (output,)

@ETK_pytorch_base
class Activation:
    """
    uses values in torch.nn.modules.activation to create an activation function
    """
    @classmethod
    def INPUT_TYPES(cls):
        activation_str_names = [i for i in torch.nn.modules.activation.__dict__.keys() if i[0] != "_"]
        return {
            "required": {
                "model": ("TORCH_MODEL",),
                "activation": (activation_str_names,),
            }
        }
    RETURN_TYPES = ("TORCH_MODEL",)
    FUNCTION = "activation"
    def activation(self, model, activation):
        act = getattr(torch.nn.modules.activation, activation)()
        model = model.insert(len(model), act)
        return (model,)

@ETK_pytorch_base
class AddReshapeLayer:
    """
    Adds a reshape layer to the model.
    """
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": ("TORCH_MODEL",),
                "shape": ("STRING", {"default": "1, -1"}),
            }
        }
    RETURN_TYPES = ("TORCH_MODEL",)
    FUNCTION = "add_reshape_layer"
    CATEGORY = "model"
    def add_reshape_layer(self, model, shape):
        import torch.nn as nn
        t = tuple([int(i) for i in shape.split(",")])
        reshape_layer_f = nn.Flatten(0)
        reshape_layer_u = nn.Unflatten(0, t)
        if isinstance(model, nn.Sequential):
            model.insert(len(model), reshape_layer_f)
            model.insert(len(model), reshape_layer_u)
        else:
            raise TypeError("The provided model is not an nn.Sequential model.")
        return (model,)

@ETK_pytorch_base
class AddTransformerLayer:
    """
    Adds a transformer encoder layer to the model.
    """
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": ("TORCH_MODEL",),
                "input_features": ("INT", {"min": 1, "max": 2 ** 24}),
                "num_heads": ("INT", {"min": 1}),
                "feedforward_dim": ("INT", {"min": 1}),
                "dropout": ("FLOAT", {"min": 0.0, "max": 1.0, "default": 0.1}),
            }
        }
    RETURN_TYPES = ("TORCH_MODEL",)
    FUNCTION = "add_transformer_layer"
    CATEGORY = "model"
    def add_transformer_layer(self, model, input_features, num_heads, feedforward_dim, dropout=0.1):
        import torch.nn as nn
        transformer_layer = nn.TransformerEncoderLayer(
            d_model=input_features,
            nhead=num_heads,
            dim_feedforward=feedforward_dim,
            dropout=dropout
        )
        if isinstance(model, nn.Sequential):
            model.insert(len(model), transformer_layer)
        else:
            raise TypeError("The provided model is not a nn.Sequential model.")
        return (model,)

@ETK_pytorch_base
class PyTorchToDevice:
    """
    moves a pytorch object to a device
    """
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "device": (["cuda", "cpu"], {"default": "cpu"}),
            },
            "optional": {
                "model": ("TORCH_MODEL",),
                "tensor": ("TORCH_TENSOR",),
            }
        }
    RETURN_TYPES = ("TORCH_MODEL", "TORCH_TENSOR",)
    FUNCTION = "to_device"
    def to_device(self, device, model=None, tensor=None):
        device_to_use = torch.device(device)
        model_out = model.to(device_to_use) if model is not None else None
        tensor_out = tensor.to(device_to_use) if tensor is not None else None
        return (model_out, tensor_out,)

@ETK_pytorch_base
class ExtractLayersAsModel:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model": ("TORCH_MODEL",),
                "start_idx": ("INT", {"min": 0}),
                "end_idx": ("INT", {"min": 0}),
                "freeze": (["True", "False"], {"default": "False"}),
            }
        }
    RETURN_TYPES = ("TORCH_MODEL", "STRING",)
    RETURN_NAMES = ("model", "shapes_str",)
    FUNCTION = "extract_layers"
    CATEGORY = "model"
    def extract_layers(self, model, start_idx, end_idx, freeze=False):
        if freeze == "True":
            freeze = True
        else:
            freeze = False
        if not isinstance(model, nn.Sequential):
            raise TypeError("The provided model is not a nn.Sequential model.")
        if start_idx < 0 or end_idx < start_idx or end_idx >= len(model):
            raise ValueError("Invalid start or end index.")
        extracted_layers = nn.Sequential(*list(model.children())[start_idx:end_idx + 1])
        shapes = []
        for i in range(start_idx, end_idx + 1):
            if hasattr(model[i], "in_features"):
                shapes.append(model[i].in_features)
            if hasattr(model[i], "out_features"):
                shapes.append(model[i].out_features)
        shapes_str = str(shapes)
        if freeze:
            for param in extracted_layers.parameters():
                param.requires_grad = False
        return (extracted_layers, shapes_str,)

@ETK_pytorch_base
class AddModelAsLayer:
    """
    given a model main and model addition adds the model addition to the end of the model main
    """
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model_main": ("TORCH_MODEL",),
                "model_addition": ("TORCH_MODEL",),
            }
        }
    RETURN_TYPES = ("TORCH_MODEL",)
    FUNCTION = "add_model"
    CATEGORY = "model"
    def add_model(self, model_main, model_addition):
        if not isinstance(model_main, nn.Sequential):
            raise TypeError("model_main is not a nn.Sequential model.")
        if not isinstance(model_addition, nn.Sequential):
            raise TypeError("model_addition is not a nn.Sequential model.")
        for layer in model_addition.children():
            model_main.add_module(f"added_layer_{len(model_main)}", layer)
        return (model_main,)

@ETK_pytorch_base
class RandomTensor:
    """
    creates a random tensor of a given shape and dtype and initialization method
    """
    @classmethod
    def INPUT_TYPES(s):
        valid_dtypes_str = [str(dt) for k, dt in torch.__dict__.items() if isinstance(dt, torch.dtype)]
        return {
            "required": {
                "shape": ("STRING", {"default": "(1,1)"}),
            },
            "optional": {
                "dtype": (valid_dtypes_str,),
                "init_method": (
                    ["rand", "randn", "randint", "randint_like", "rand_like", "randn_like"], {"default": "rand"}),
            }
        }
    RETURN_TYPES = ("TORCH_TENSOR",)
    FUNCTION = "random_tensor"
    CATEGORY = "tensor"
    def random_tensor(self, shape, dtype=str(torch.float32), init_method="rand"):
        valid_dtypes_str = [str(dt) for k, dt in torch.__dict__.items() if isinstance(dt, torch.dtype)]
        valid_dtypes = [dt for k, dt in torch.__dict__.items() if isinstance(dt, torch.dtype)]
        if dtype not in valid_dtypes_str:
            raise ValueError(f"Unsupported dtype: {dtype}")
        dto = valid_dtypes[valid_dtypes_str.index(dtype)]

        shape = shape[1:-1]
        shape = [int(s) for s in shape.split(",")]
        shape = tuple(shape)

        if init_method == "rand":
            output = torch.rand(shape, dtype=dto)
        elif init_method == "randn":
            output = torch.randn(shape, dtype=dto)
        elif init_method == "randint":
            output = torch.randint(0, 100, shape, dtype=dto)
        elif init_method == "randint_like":
            temp = torch.empty(shape, dtype=dto)
            output = torch.randint_like(temp, 0, 100)
        elif init_method == "rand_like":
            temp = torch.empty(shape, dtype=dto)
            output = torch.rand_like(temp)
        elif init_method == "randn_like":
            temp = torch.empty(shape, dtype=dto)
            output = torch.randn_like(temp)
        else:
            raise ValueError(f"Unsupported init_method: {init_method}")
        return (output,)

@ETK_pytorch_base
class GridSearchTraining:
    """
    Performs grid search training using the TrainModel class.
    """
    @classmethod
    def INPUT_TYPES(cls):
        default_param_grid = json.dumps({
            'epochs': [1, 2],
            'batch_size': [32, 64],
            'loss_function': ['MSELoss', 'CrossEntropyLoss']
        })
        return {
            "required": {
                "model": ("TORCH_MODEL",),
                "param_grid": ("STRING", {"default": default_param_grid}),
            },
            "optional": {
                "dataset": ("TORCH_DATASET",),
                "features tensor": ("TORCH_TENSOR",),
                "labels tensor": ("TORCH_TENSOR",),
            }
        }
    RETURN_TYPES = ("TORCH_MODEL", "DICT", "LIST")
    RETURN_NAMES = ("best_model", "best_params", "all_metrics")
    FUNCTION = "grid_search_train"
    def grid_search_train(self, model, param_grid, **kwargs):
        from itertools import product
        param_grid_dict = json.loads(param_grid)
        best_model = None
        best_params = None
        all_metrics = []
        param_combinations = list(product(*param_grid_dict.values()))
        for params in param_combinations:
            train_kwargs = dict(zip(param_grid_dict.keys(), params))
            train_kwargs.update(kwargs)
            trainer = TrainModel()
            trained_model, metrics, _ = trainer.train(model, **train_kwargs)
            all_metrics.append(metrics)
            if best_model is None or self.is_better(metrics, all_metrics):
                best_model = trained_model
                best_params = train_kwargs
        return best_model, best_params, all_metrics
    @staticmethod
    def is_better(current_metrics, all_metrics):
        return current_metrics[-1] < min(all_metrics, key=lambda m: m[-1])[-1]

@ETK_pytorch_base
class SaveTorchTensor:
    """
    Saves a torch tensor to a file using torch.save which has the required parameters of obj and f
    """
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "tensor": ("TORCH_TENSOR", {"default": None}),
                "file": ("STRING", {"default": "/somewhere/some.pt"}),
            },
        }
    RETURN_TYPES = ("STRING",)
    FUNCTION = "save_torch_tensor"
    OUTPUT_NODE = True
    def save_torch_tensor(self, tensor, file):
        torch.save(tensor, file)
        return (file,)

@ETK_pytorch_base
class LoadTorchTensor:
    """
    Loads a torch tensor from a file using torch.load
    """
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "file": ("STRING", {"default": "/somewhere/some.pt"}),
            },
        }
    RETURN_TYPES = ("TORCH_TENSOR",)
    FUNCTION = "load_torch_tensor"
    def load_torch_tensor(self, file):
        return (torch.load(file),)

@ETK_pytorch_base
class FuncModifyModel:
    """
    allows users to modify a model using a function
    """
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": ("TORCH_MODEL",),
                "function": ("STRING", {"default": "model", "multiline": True}),
            }
        }
    RETURN_TYPES = ("TORCH_MODEL",)
    FUNCTION = "modify_model"
    def modify_model(self, model, function):
        from copy import deepcopy
        import re

        old_model = deepcopy(model)

        # Safe function execution - only allow specific patterns
        safe_functions = {
            r"model\.add_module\('(\w+)', nn\.(\w+)\(\)\)": lambda m: model.add_module(m.group(1), getattr(torch.nn, m.group(2))()),
            r"model\.add_module\('(\w+)', nn\.(\w+)\(([^)]+)\)\)": lambda m: model.add_module(m.group(1), getattr(torch.nn, m.group(2))()),
            r"model\.eval\(\)": lambda m: model.eval(),
            r"model\.train\(\)": lambda m: model.train(),
        }

        executed = False
        for pattern, func in safe_functions.items():
            match = re.match(pattern, function.strip())
            if match:
                try:
                    func(match)
                    executed = True
                    break
                except Exception:
                    continue

        if not executed:
            # For safety, just add a ReLU layer as default modification
            model.add_module('safe_relu', torch.nn.ReLU())

        new_model = model
        if str(old_model) == str(new_model):  # Compare string representations
            raise ValueError("The model was not modified.")
        return (new_model,)

@ETK_pytorch_base
class PlotSeriesString:
    """
    Plots a series of strings
    """
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "series": ("STRING", {"default": "1\n2\n3\n4\n5\n6\n7\n8\n9\n10"}),
            },
            "optional": {
                "title": ("STRING", {"default": "Series Plot"}),
                "xlabel": ("STRING", {"default": "X-axis"}),
                "ylabel": ("STRING", {"default": "Y-axis"}),
            }
        }
    RETURN_TYPES = ("IMAGE",)
    FUNCTION = "plot_series_string"
    OUTPUT_NODE = True
    def plot_series_string(self, series, title="Series Plot", xlabel="X-axis", ylabel="Y-axis"):
        import matplotlib.pyplot as plt
        import numpy as np
        import io
        import PIL
        from PIL import Image
        if isinstance(series, str):
            series = series.split("\n")
        elif isinstance(series, list):
            pass
        else:
            raise ValueError("series must be a string or a list")
        series = [float(s) for s in series]
        plt.plot(series)
        plt.title(title)
        plt.xlabel(xlabel)
        plt.ylabel(ylabel)
        buf = io.BytesIO()
        plt.savefig(buf, format='png')
        buf.seek(0)
        im = Image.open(buf)
        im = np.array(im)
        im = torch.tensor(im)
        im = im.unsqueeze(0)
        im = im.float()
        im = im / 255.0
        plt.close()
        buf.close()
        return (im,)


@ETK_pytorch_base
class AddBatchNormLayer:
    """Adds a BatchNorm layer to an ``nn.Sequential`` model."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": ("TORCH_MODEL",),
                "num_features": ("INT", {"min": 1}),
                "one_d": ("BOOLEAN", {"default": True}),
            }
        }

    RETURN_TYPES = ("TORCH_MODEL",)
    FUNCTION = "add_batchnorm_layer"
    CATEGORY = "model"

    def add_batchnorm_layer(self, model, num_features, one_d=True):
        layer = nn.BatchNorm1d(num_features) if one_d else nn.BatchNorm2d(num_features)
        if isinstance(model, nn.Sequential):
            model.insert(len(model), layer)
        else:
            raise TypeError("The provided model is not a nn.Sequential model.")
        return (model,)


@ETK_pytorch_base
class AddDropoutLayer:
    """Adds a Dropout layer to an ``nn.Sequential`` model."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": ("TORCH_MODEL",),
                "p": ("FLOAT", {"min": 0.0, "max": 1.0, "default": 0.5}),
            },
            "optional": {
                "inplace": (["True", "False"], {"default": "False"}),
            },
        }

    RETURN_TYPES = ("TORCH_MODEL",)
    FUNCTION = "add_dropout_layer"
    CATEGORY = "model"

    def add_dropout_layer(self, model, p=0.5, inplace="False"):
        inplace = inplace == "True"
        layer = nn.Dropout(p=p, inplace=inplace)
        if isinstance(model, nn.Sequential):
            model.insert(len(model), layer)
        else:
            raise TypeError("The provided model is not a nn.Sequential model.")
        return (model,)


@ETK_pytorch_base
class SetModelTrainable:
    """Sets ``requires_grad`` for all parameters of a model."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "model": ("TORCH_MODEL",),
                "requires_grad": (["True", "False"], {"default": "False"}),
            }
        }

    RETURN_TYPES = ("TORCH_MODEL",)
    FUNCTION = "set_trainable"
    CATEGORY = "model"

    def set_trainable(self, model, requires_grad="False"):
        flag = True if requires_grad in (True, "True") else False
        for param in model.parameters():
            param.requires_grad = flag
        return (model,)