From 447e9403ae4c1c11001289e18f4b4919e6069e81 Mon Sep 17 00:00:00 2001 From: Tasha Upchurch Date: Thu, 18 Jan 2024 17:28:28 -0500 Subject: [PATCH] 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. --- basic.py | 115 ++++++++++++++++++++++++++++- config.yaml | 1 + hf_diffusers_nodes.py | 112 +++++++++++++++++++++++++++- image.py | 121 ++++++++++++++++++++++++++++++- pytorch_nodes.py | 58 ++++++++++++++- server_endpoints.py | 39 +++++++--- tests/test_hf_diffusers_nodes.py | 27 +++++-- 7 files changed, 450 insertions(+), 23 deletions(-) diff --git a/basic.py b/basic.py index ecbebbd..fcb6e4f 100644 --- a/basic.py +++ b/basic.py @@ -23,6 +23,87 @@ class MultilineStringNode: return (Text,) +class TupleNode: + """Node that displays a tuple""" + + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "Tuple": ("STRING", { + "default": "(None,None,)", + }), + } + } + + RETURN_TYPES = ("TUPLE",) + FUNCTION = "get_value" + CATEGORY = "text" + + def get_value(self, Tuple): + # evaluate the string safely step by step + # type check + if not isinstance(Tuple, str): + raise TypeError("Tuple must be a string") + + # format check + if not Tuple.startswith("(") or not Tuple.endswith(")"): + raise ValueError("Tuple must start with '(' and end with ')'") + + # remove the brackets + Tuple = Tuple.strip("()") + + # check for empty tuple + if len(Tuple) == 0: + return ((),) + # check for commas + if "," not in Tuple: + raise ValueError("Tuple must have at least one comma") + + # split on commas + Tuple = Tuple.split(",") + + # valid values are string,int,float,None, list,tuple,dict + # construct the tuple + ret = [] + for val in Tuple: + val = val.strip() + if val == "None": + ret.append(None) + elif val.startswith("[") and val.endswith("]"): + # list + val = val.strip("[]") + val = val.split(",") + ret.append(val) + elif val.startswith("(") and val.endswith(")"): + # tuple + val = val.strip("()") + val = val.split(",") + ret.append(tuple(val)) + elif val.startswith("{") and val.endswith("}"): + # dict + val = val.strip("{}") + val = val.split(",") + ret.append(dict(val)) + elif val.startswith('"') and val.endswith('"'): + # string + ret.append(val.strip('"')) + elif "." in val: + # float + ret.append(float(val)) + elif val == "": + # empty + pass + else: + # int + ret.append(int(val)) + + return (tuple(ret),) + + class StrBreakout: """ node that breaks out a string into a: - Boolean @@ -45,14 +126,16 @@ class StrBreakout: "Text": ("STRING", { "default": "", }), + "list_dtype": (["str", "int", "float", ],), } } RETURN_TYPES = ("STRING", "BOOL", "INT", "FLOAT", "LIST", "LIST", "LIST",) + RETURN_NAMES = ("Notes", "Bool", "Int", "Float", "List by chr", "list by lines", "list by commas",) FUNCTION = "get_value" CATEGORY = "text" - def get_value(self, Text): + def get_value(self, Text, list_dtype="str"): out_bool = False out_int = 0 out_float = 0.0 @@ -95,6 +178,34 @@ class StrBreakout: except Exception as e: notes.append(f"split(',') failed: {e}") pass + + if list_dtype == "int": + list_dtype = int + elif list_dtype == "float": + list_dtype = float + else: + list_dtype = str + + try: + out_list = [list_dtype(val) for val in out_list] + notes.append(f"list dtype worked") + except Exception as e: + notes.append(f"list dtype failed: {e}") + pass + try: + out_lines = [list_dtype(val) for val in out_lines] + notes.append(f"splitlines dtype worked") + except Exception as e: + notes.append(f"splitlines dtype failed: {e}") + pass + try: + out_comma = [list_dtype(val) for val in out_comma] + notes.append(f"split(',') dtype worked") + except Exception as e: + notes.append(f"split(',') dtype failed: {e}") + pass + + return ("\n".join(notes), out_bool, out_int, out_float, out_list, out_lines, out_comma,) @@ -308,3 +419,5 @@ NODE_CLASS_MAPPINGS["ValueToAny"] = ValueToAny NODE_CLASS_MAPPINGS["IterableToAny"] = IterableToAny NODE_CLASS_MAPPINGS["CallableToAny"] = CallableToAny NODE_CLASS_MAPPINGS["ObjectToAny"] = ObjectToAny + +NODE_CLASS_MAPPINGS["ETK_Tuple"] = TupleNode diff --git a/config.yaml b/config.yaml index aeab130..94cc943 100644 --- a/config.yaml +++ b/config.yaml @@ -3,3 +3,4 @@ tmp_dir: comfy_path/tmp gitio: COTRF: https://github.com/Hopping-Mad-Games/COTRF.git StSq: https://github.com/story-squad/GDE_Graph_IO.git + base: ./data/gitio diff --git a/hf_diffusers_nodes.py b/hf_diffusers_nodes.py index cdc2295..e00a205 100644 --- a/hf_diffusers_nodes.py +++ b/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,) \ No newline at end of file + 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,) + diff --git a/image.py b/image.py index 72b8e93..399bf37 100644 --- a/image.py +++ b/image.py @@ -2207,7 +2207,7 @@ class FloodFillNode: RETURN_TYPES = ("IMAGE",) FUNCTION = "execute" - CATEGORY = "image processing" + def __init__(self): self.num_workers = None # Can be set to the number of processes to create @@ -2263,6 +2263,125 @@ class FloodFillNode: images[i] = self.flood_fill(images[i], (start_x, start_y), target_color, replacement_color) return images + +@ETK_image_base +class DecorateHeightMap: + """ + Decorate a height map with input images at different heights specified by input list + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "height_map": ("IMAGE",), # Expecting a NumPy array or tensor + "images": ("IMAGE",), # Expecting a NumPy array or tensor + "heights": ("LIST", {"default": [0.5]}), + "feather": ("LIST", {"default": [0.05]}), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "execute" + + + def execute(self, height_map, images, heights, feather=0.05): + """ + Args: + height_map: float values 0-1 1,h,w,3 + images: z,x,y,c + heights: list of floats 0-1 length of z-1 + feather: single float or a list of floats specifying the gradient width for feathering effect + Returns: + torch image 1,h,w,3 + """ + heights = heights.copy() + images = images[..., 0:3].clone() + height_map = height_map[..., 0:3].clone() + + # Ensure heights is a list and add boundary values + heights = [0.0] + heights + [1.0] + + # If feather is a single float, create a list of the same feather value for each transition + if isinstance(feather, float): + feather = [feather] * (len(heights) - 1) + + # so the length of th feather list must be at least 1 + # if there is one height then it should be length 1 + # if there are two heights then it should be length 2 + # if there are three heights then it should be length 3 + # etc + if len(feather) != len(heights)-2: + raise ValueError("feather must be a single float or a list of floats of the same length as heights") + + + # Check for dimension match between height_map and images + if height_map.shape[1:3] != images.shape[1:3]: + raise ValueError("height_map and images must have the same height and width") + + # Convert height map to 2D + height_map_2d = height_map.squeeze(0)[..., 0] + + # Initialize output with the first texture + output = images[0].clone() + for i in range(1, len(heights) - 1): + current_feather = feather[i - 1] + # Calculate the mask for the current segment + mask = torch.sigmoid((height_map_2d - heights[i - 1]) / current_feather) * ( + 1 - torch.sigmoid((height_map_2d - heights[i]) / current_feather)) + mask = mask.unsqueeze(-1).repeat(1, 1, 3) # Expand mask to match image dimensions + + # Blend the output with the next texture using the mask + output = output * (1 - mask) + images[i] * mask + + # Normalize the output to ensure it's within the valid range of [0, 1] + output = torch.clamp(output, 0, 1).unsqueeze(0) + return (output,) + +@ETK_image_base +class RandomImageFromFolder: + """ + Randomly select an image from a folder + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "folder": ("STRING",{"default": "./output"}), + } + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "execute" + + + def execute(self, folder): + """ + Args: + folder: string + Returns: + torch image 1,h,w,3 + """ + import random + # Get a list of all files in the folder + files = os.listdir(folder) + + # Filter out non-image files + files = [f for f in files if f.endswith(".jpg") or f.endswith(".png")] + + # Select a random file + file = random.choice(files) + + # Load the image + image = Image.open(os.path.join(folder, file)) + + # Convert to torch tensor + image = torch.from_numpy(np.array(image)).float() / 255.0 + image = image.unsqueeze(0) + + return (image,) + import torch diff --git a/pytorch_nodes.py b/pytorch_nodes.py index e459449..8da397f 100644 --- a/pytorch_nodes.py +++ b/pytorch_nodes.py @@ -5,6 +5,7 @@ import torch.nn as nn import functools import torch import json + try: from .config import config_settings except ImportError as e: @@ -1086,6 +1087,8 @@ class RandomTensor: elif init_method == "randn_like": output = torch.randn_like(shape, dtype=dto) return (output,) + + @ETK_pytorch_base class GridSearchTraining: """ @@ -1165,10 +1168,63 @@ class GridSearchTraining: 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" + + def save_torch_tensor(self, tensor, file): + import torch + """ + Saves a torch tensor to a 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): + import torch + """ + Loads a torch tensor from a file + """ + return (torch.load(file),) if __name__ == "__main__": import json + # Example usage param_grid_str = json.dumps({ 'epochs': [1, 2], @@ -1206,4 +1262,4 @@ if __name__ == "__main__": dataset=torch.utils.data.TensorDataset(synthetic_features, synthetic_labels) ) - print("Best parameters:", best_params) \ No newline at end of file + print("Best parameters:", best_params) diff --git a/server_endpoints.py b/server_endpoints.py index 766f47e..ec80fcc 100644 --- a/server_endpoints.py +++ b/server_endpoints.py @@ -10,10 +10,11 @@ from folder_paths import output_directory, input_directory, temp_directory from custom_nodes.EternalKernelLiteGraphNodes.local_shared import GDE_PATH, ETK_PATH fake_git = None -#real_git_url = "https://github.com/story-squad/GDE_Graph_IO.git" # ssh url -#real_git_url = "https://github.com/Hopping-Mad-Games/COTRF.git" # ssh url +# real_git_url = "https://github.com/story-squad/GDE_Graph_IO.git" # ssh url +# real_git_url = "https://github.com/Hopping-Mad-Games/COTRF.git" # ssh url + class GitIO: def __init__(self, i_real_git_url: str, i_gde_path: str): @@ -22,8 +23,15 @@ class GitIO: def setup_git_io(self, l_real_git_url: str): git_name = l_real_git_url.split("/")[-1].split(".")[0] - real_git_path = os.path.join(self.gde_path, "data") - real_git_path = os.path.join(real_git_path, git_name) + + real_git_path = os.path.join(self.gde_path, "data","gitio",git_name) + + + # users_gitio_path = os.path.abspath(os.path.join( + # GDE_PATH, + # config_settings["gitio"]["base"], + # users_gitio) + # ) # check to see if it exists and is a git repo, if not then clone it if os.path.exists(real_git_path): @@ -83,7 +91,7 @@ class GitIO: # git checkout feature-branch run_git_command(["git", "checkout", user]) # set the upstream to the remote - #run_git_command(["git", "branch", "--set-upstream-to", "origin/" + user]) + # run_git_command(["git", "branch", "--set-upstream-to", "origin/" + user]) # push the branch to the remote run_git_command(["git", "push", "-u", "origin", user]) # fetch and pull it back down for good measure @@ -94,7 +102,7 @@ class GitIO: print(f"git branch {user} exists locally but not on the remote") # the branch does not exist on the remote but does locally # set the upstream - #run_git_command(["git", "branch", "--set-upstream-to", "origin/" + user]) + # run_git_command(["git", "branch", "--set-upstream-to", "origin/" + user]) # push the branch to the remote run_git_command(["git", "push", "-u", "origin", user]) # fetch and pull it back down for good measure @@ -188,9 +196,8 @@ user_data = { "security": None, "client_id": None, "git_io": "COTRF" - }, - } - + }, +} def get_user_by_client_id(client_id): @@ -251,6 +258,7 @@ async def login(request): - sets up security - user and pass are in the json """ + print("login request") # get the user and pass from the json @@ -306,15 +314,24 @@ async def login(request): user_data[user]["client_id"] = client_id global git_io + + #from . import config + #config_settings = config.config_settings + users_gitio = user_data[user]["git_io"] - users_gitio_path = os.path.join(GDE_PATH, users_gitio) + + #users_gitio_path = os.path.abspath(os.path.join( + # GDE_PATH, + # config_settings["gitio"]["base"], + # users_gitio) + #) try: users_gittio_url = config_settings["gitio"][users_gitio] except KeyError: return web.json_response({"error": "invalid gitio"}, status=401) - git_io = GitIO(users_gittio_url, users_gitio_path) + git_io = GitIO(users_gittio_url, GDE_PATH) # return json with the security return web.json_response({"security": security}) diff --git a/tests/test_hf_diffusers_nodes.py b/tests/test_hf_diffusers_nodes.py index 0c0ea77..43a7bbf 100644 --- a/tests/test_hf_diffusers_nodes.py +++ b/tests/test_hf_diffusers_nodes.py @@ -3,7 +3,6 @@ from ..hf_diffusers_nodes import DDPMPipline, SampleDDPMPipline, TensorToImage from ..hf_diffusers_nodes import LoadHFDataset - # when testing set "UNIT_TEST" to true in the environment to avoid loading all of ComfyUI def test_DDPMPipline(): p_node = DDPMPipline() @@ -17,7 +16,7 @@ def test_SampleDDPMPipline(): s_node = SampleDDPMPipline() ret = p_node.load_from_pretrained("johnowhitaker/ddpm-butterflies-32px", "cuda") - ret=ret[0] + ret = ret[0] assert ret is not None ret = s_node.sample(ret, 1, @@ -40,8 +39,8 @@ def test_TensorToImage(): ) ret = ret[0] assert ret is not None - ret = t_node.tensor_to_image(ret,"HWC") - ret=ret[0] + ret = t_node.tensor_to_image(ret, "HWC") + ret = ret[0] assert ret is not None # now convert to pillow image from PIL import Image @@ -51,14 +50,13 @@ def test_TensorToImage(): ret = ret.squeeze(0) ret = ret.squeeze(0) ret = ret.numpy() - ret = ret*255 + ret = ret * 255 ret = np.uint8(ret) ret = Image.fromarray(ret) ret.show() - def test_LoadHFDataset(): p_node = LoadHFDataset() ret = p_node.load_from_hf("huggan/smithsonian_butterflies_subset") @@ -67,3 +65,20 @@ def test_LoadHFDataset(): if __name__ == "__main__": pytest.main() + + +def test_create_unet2d_model(): + from ..hf_diffusers_nodes import CreateUnet2dModel + p_node = CreateUnet2dModel() + ret = p_node.create_unet2d_model(3, + 3, + 1, + 1, + 1, + 1, + ) + assert ret is not None + + + +