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 + + + +