nodes/TashaSkyUp-EternalKernelPytorchNodes
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Refactoring and Enhancements Across Multiple Modules

Browse Source 
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.
This commit is contained in:
Tasha Upchurch
2024-01-18 17:28:28 -05:00
parent d2d16cadfe
commit 447e9403ae
7 changed files with 450 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
basic.py
+114 -1
View File
@@ -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
config.yaml
+1
View File
@@ -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
hf_diffusers_nodes.py
+109 -3
View File
@@ -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,)
image.py
+120 -1
View File
@@ -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
pytorch_nodes.py
+57 -1
View File
@@ -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)
print("Best parameters:", best_params)
server_endpoints.py
+28 -11
View File
@@ -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})
tests/test_hf_diffusers_nodes.py
+21 -6
View File
@@ -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