Refactoring and Enhancements Across Multiple Modules

This commit introduces a series of refactoring and enhancements across multiple modules, improving code readability, consistency, and adding new functionalities.

Highlights:
1. `tests/test_hf_diffusers_nodes.py`:
   - Removed unnecessary whitespace.
   - Improved variable assignment consistency (e.g., `ret=ret[0]` to `ret = ret[0]`).
   - Added a new test case `test_create_unet2d_model` for Unet2d model creation.

2. `server_endpoints.py`:
   - Cleaned up commented-out code and added more descriptive comments.
   - Refactored the construction of `real_git_path` for better readability and future maintenance.
   - Implemented minor formatting improvements for better code structure.

3. `basic.py`:
   - Introduced a new class `TupleNode` to handle tuple data types.
   - Modified `StrBreakout` to include dynamic typing for list conversion.
   - Enhanced code readability with better variable naming and added comments for clarity.

4. `pytorch_nodes.py`:
   - Added import statements within functions for `torch`.
   - New classes `SaveTorchTensor` and `LoadTorchTensor` added for saving and loading torch tensors.
   - `GridSearchTraining` class updated with minor formatting changes.

5. `image.py`:
   - In `FloodFillNode`, removed a redundant category assignment.
   - Added two new classes: `DecorateHeightMap` and `RandomImageFromFolder`. These classes provide functionalities for decorating height maps with images and selecting random images from a folder, respectively.

6. `hf_diffusers_nodes.py`:
   - Fixed formatting in tensor format options and added a new line for clarity.
   - New class `CreateUnet2dModel` added for creating a UNet2D model with customizable parameters.
   - Introduced `TrainUnet2dModel` for training a UNet2D model with a dataset, including detailed setup for the training process.

7. `config.yaml`:
   - Added a new base path for `gitio` to centralize the configuration.

These changes aim to improve the maintainability of the codebase, streamline testing processes, and expand the functionalities within the modules, particularly related to image processing and model training in PyTorch and diffusers.

hf_diffusers_nodes.py

@@ -77,7 +77,7 @@ class SampleDDPMPipline:
             },
         }
 
-    RETURN_TYPES = ("TENSOR",)
+    RETURN_TYPES = ("TORCH_TENSOR",)
     FUNCTION = "sample"
 
     def sample(
@@ -105,7 +105,7 @@ class TensorToImage:
         return {
             "required": {
                 "tensor": ("TORCH_TENSOR", {}),
-                "incoming_format": (["BHWC", "BCHW","HWC"], {"default": "BHWC"}
+                "incoming_format": (["BHWC", "BCHW", "HWC"], {"default": "BHWC"}
                                     ),
             },
         }
@@ -124,6 +124,7 @@ class TensorToImage:
 
         return (tensor,)
 
+
 @ETK_HF_Diffusers_base
 class LoadHFDataset:
     """A node that loads the specified HFDataset"""
@@ -144,4 +145,109 @@ class LoadHFDataset:
         ret = load_dataset(dataset_name,
                            split="train",
                            )
-        return (ret,)
+        return (ret,)
+
+
+@ETK_HF_Diffusers_base
+class CreateUnet2dModel:
+    """A node that creates a unet2d model"""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        down_blocks = ("DownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D")
+        mid_blocks = ('UNetMidBlock2D', 'UnCLIPUNetMidBlock2D')
+        up_blocks = ("AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D", "UpBlock2D")
+
+        return {
+            "required": {
+                "name": ("STRING", {"default": "UNet2D"}),
+            },
+            "optional": {
+                "sample_size": ("TUPLE", {"default": (32, 32, 3,)}),
+                "in_channels": ("INT", {"default": 3}),
+                "out_channels": ("INT", {"default": 3}),
+                "norm_num_groups": ("INT", {"default": 32}),
+                "layers_per_block": ("INT", {"default": 2}),
+                "down_block_types": ("TUPLE", {"default": (down_blocks[0],)}),
+                "mid_block_scale_factor": ("FLOAT", {"default": 1.0}),
+                "up_block_types": ("TUPLE", {"default": (up_blocks[0],)}),
+                "block_out_channels": ("TUPLE", {"default": (64, 128, 256, 512,)}),
+            },
+        }
+
+    RETURN_TYPES = ("TORCH_MODEL",)
+    FUNCTION = "create_unet2d_model"
+
+    def create_unet2d_model(self, name, sample_size, in_channels, out_channels, norm_num_groups, layers_per_block,
+                            down_block_types=("AttnDownBlock2D",),
+                            mid_block_scale_factor=("UNetMidBlock2D",),
+                            up_block_types=("AttnUpBlock2D",),
+                            block_out_channels=(64, 128, 256, 512,)
+                            ):
+        from diffusers import UNet2DModel
+        with torch.inference_mode(False):
+            ret = UNet2DModel(
+                sample_size=sample_size,
+                in_channels=in_channels,
+                out_channels=out_channels,
+                norm_num_groups=norm_num_groups,
+                layers_per_block=layers_per_block,
+                down_block_types=down_block_types,
+                mid_block_scale_factor=mid_block_scale_factor,
+                up_block_types=up_block_types,
+                block_out_channels=block_out_channels,
+
+            )
+        for param in ret.parameters():
+            param.requires_grad = True
+            #param.retains_grad = True
+
+        return (ret,)
+
+@ETK_HF_Diffusers_base
+class TrainUnet2dModel:
+    """A node that trains a unet2d model"""
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model": ("TORCH_MODEL", {}),
+                "dataset": ("HF_DATASET", {}),
+                "batch_size": ("INT", {"default": 1}),
+                "num_epochs": ("INT", {"default": 1}),
+                "learning_rate": ("FLOAT", {"default": 1e-4}),
+            },
+        }
+
+    RETURN_TYPES = ("TORCH_MODEL",)
+    FUNCTION = "train_unet2d_model"
+
+    def train_unet2d_model(self, model, dataset, batch_size, num_epochs, learning_rate):
+        from torch.utils.data import DataLoader
+        # import mse loss
+        from torch.nn import MSELoss
+        from torch.optim import Adam
+        from diffusers import UNet2DModel
+
+
+        model = model.train()
+        model = model.cuda()
+
+        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
+
+        optimizer = Adam(model.parameters(), lr=learning_rate)
+        loss_fn = MSELoss()
+
+        for epoch in range(num_epochs):
+            for batch in dataloader:
+                batch = batch["image"]
+                batch = batch.cuda()
+                optimizer.zero_grad()
+                loss = loss_fn(model, batch)
+                loss.backward()
+                optimizer.step()
+                print(loss.item())
+
+        return (model,)
+