Tasha Upchurch
b57701a669
Simplified class imports and added fallback mechanisms for missing modules. Enhanced image.py with new nodes like LoadImageFromPath and SaveImageToPath. Updated server_endpoints.py with file-backed dictionary for user data persistence. Refined node creation and sampling logic in pytorch_nodes.py and hf_diffusers_nodes.py for improved modularity and consistency."
279 lines
8.3 KiB
Python
279 lines
8.3 KiB
Python
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
|
|
|
|
#NODE_CLASS_MAPPINGS = {}
|
|
#NODE_DISPLAY_NAME_MAPPINGS = {}
|
|
from . import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
|
|
|
|
|
|
def ETK_HF_Diffusers_base(cls):
|
|
cls.CATEGORY = "ETK/HF/Diffusers"
|
|
# 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_HF_Diffusers_base
|
|
class DDPMPipline:
|
|
"""A node that loads the specified DDPMPipeline, supports only from_pretrained for now"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model_name": ("STRING", {"default": "johnowhitaker/ddpm-butterflies-32px"}),
|
|
"device": (["cpu", "cuda"], {"default": "cuda"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("DDPM_PL",)
|
|
FUNCTION = "load_from_pretrained"
|
|
|
|
def load_from_pretrained(self, model_name, device):
|
|
from diffusers import DDPMPipeline
|
|
ret = DDPMPipeline.from_pretrained(model_name, device=device)
|
|
return (ret,)
|
|
|
|
|
|
@ETK_HF_Diffusers_base
|
|
class SampleDDPMPipline:
|
|
"""A node that samples from the specified DDPMPipeline"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": ("DDPM_PL", {}),
|
|
"batch_size": ("INT", {"default": 1}),
|
|
"num_timesteps": ("INT", {"default": 1000}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("TORCH_TENSOR",)
|
|
FUNCTION = "sample"
|
|
|
|
def sample(
|
|
self,
|
|
model,
|
|
batch_size,
|
|
num_timesteps,
|
|
):
|
|
ret = model(
|
|
batch_size=batch_size,
|
|
generator=None,
|
|
num_inference_steps=num_timesteps,
|
|
output_type="tensor",
|
|
)
|
|
ret = torch.tensor(ret.images)
|
|
return (ret,)
|
|
|
|
|
|
@ETK_HF_Diffusers_base
|
|
class TensorToImage:
|
|
"""A node that converts a tensor to an image"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"tensor": ("TORCH_TENSOR", {}),
|
|
"incoming_format": (["BHWC", "BCHW", "HWC"], {"default": "BHWC"}
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "tensor_to_image"
|
|
|
|
def tensor_to_image(self, tensor, incoming_format):
|
|
# outgoing format is BHWC torch tensor float 32 0-1
|
|
|
|
if incoming_format == "BCHW":
|
|
tensor = tensor.permute(0, 2, 3, 1)
|
|
elif incoming_format == "HWC":
|
|
tensor = tensor.unsqueeze(0)
|
|
tensor = tensor.clamp(0, 1)
|
|
|
|
return (tensor,)
|
|
|
|
|
|
@ETK_HF_Diffusers_base
|
|
class LoadHFDataset:
|
|
"""A node that loads the specified HFDataset"""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"dataset_name": ("STRING", {"default": "huggan/smithsonian_butterflies_subset"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("HF_DATASET",)
|
|
FUNCTION = "load_from_hf"
|
|
|
|
def load_from_hf(self, dataset_name):
|
|
from datasets import load_dataset
|
|
ret = load_dataset(dataset_name,
|
|
split="train",
|
|
)
|
|
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,)
|
|
|
|
# @ETK_HF_Diffusers_base
|
|
# class LoadModelFromHF:
|
|
# """A node that loads the specified model from huggingface"""
|
|
#
|
|
# @classmethod
|
|
# def INPUT_TYPES(cls):
|
|
# return {
|
|
# "required": {
|
|
# "model_name": ("STRING", {"default": "johnowhitaker/ddpm-butterflies-32px"}),
|
|
# },
|
|
# }
|
|
#
|
|
# RETURN_TYPES = ("TORCH_MODEL",)
|
|
# FUNCTION = "load_model_from_hf"
|
|
#
|
|
# def load_model_from_hf(self, model_name):
|
|
# from huggingface_hub import hf_hub_url
|
|
# from diffusers import DDPMPipeline
|
|
# ret = DDPMPipeline.from_pretrained(model_name, device="cuda")
|
|
# return (ret,)
|
|
|
|
|
|
|
|
|