nodes/zeroxoxo-ComfyUI-Fast-Style-Transfer
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add files via upload

Browse Source
This commit is contained in:
zeroxoxo
2024-08-13 09:48:08 +05:00
committed by GitHub
parent 49174ffd08
commit 2f015c1c8a
4 changed files with 349 additions and 188 deletions
Download Patch File Download Diff File Expand all files Collapse all files
FastStyleTransferNode.py
+111 -59
View File
@@ -1,50 +1,22 @@
"""
This node is a simple conversion of this repository into ComfyUI ecosystem:
These nodes are a simple conversion of these repositories into ComfyUI ecosystem:
https://github.com/rrmina/fast-neural-style-pytorch.git
https://github.com/gordicaleksa/pytorch-neural-style-transfer.git
Some of the code is written by ChatGPT4-o
"""
import torch
import torch.nn as nn
from torchvision import models
import time
import os
import folder_paths
import subprocess as sp
import sys
class VGG16(nn.Module):
def __init__(self, vgg_path=os.path.join(folder_paths.base_path, "custom_nodes/ComfyUI-Fast-Style-Transfer/vgg/vgg16-00b39a1b.pth"), train=False):
super(VGG16, self).__init__()
# Load VGG Skeleton, Pretrained Weights
vgg16_features = models.vgg16(pretrained=False)
vgg16_features.load_state_dict(torch.load(vgg_path), strict=False)
self.features = vgg16_features.features
# Turn-off Gradient History (on for testing train)
for param in self.features.parameters():
param.requires_grad = train
def forward(self, x):
layers = {'3': 'relu1_2', '8': 'relu2_2', '15': 'relu3_3', '22': 'relu4_3'}
features = {}
for name, layer in self.features._modules.items():
x = layer(x)
if name in layers:
features[layers[name]] = x
if (name=='22'):
break
return features
class ConvLayer(nn.Module):
# ML classes
class ConvolutionalLayer(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size, stride, norm="instance"):
super(ConvLayer, self).__init__()
super(ConvolutionalLayer, self).__init__()
# Padding Layers
self.padding_size = kernel_size // 2
self.reflection_pad = nn.ReflectionPad2d(self.padding_size)
@@ -77,9 +49,9 @@ class ResidualLayer(nn.Module):
"""
def __init__(self, channels=128, kernel_size=3):
super(ResidualLayer, self).__init__()
self.conv1 = ConvLayer(channels, channels, kernel_size, stride=1)
self.conv1 = ConvolutionalLayer(channels, channels, kernel_size, stride=1)
self.relu = nn.ReLU()
self.conv2 = ConvLayer(channels, channels, kernel_size, stride=1)
self.conv2 = ConvolutionalLayer(channels, channels, kernel_size, stride=1)
def forward(self, x):
identity = x # preserve residual
@@ -89,9 +61,9 @@ class ResidualLayer(nn.Module):
return out
class DeconvLayer(nn.Module):
class DeconvolutionalLayer(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size, stride, output_padding, norm="instance"):
super(DeconvLayer, self).__init__()
super(DeconvolutionalLayer, self).__init__()
# Transposed Convolution
padding_size = kernel_size // 2
@@ -121,11 +93,11 @@ class TransformerNetwork(nn.Module):
def __init__(self):
super(TransformerNetwork, self).__init__()
self.ConvBlock = nn.Sequential(
ConvLayer(3, 32, 9, 1),
ConvolutionalLayer(3, 32, 9, 1),
nn.ReLU(),
ConvLayer(32, 64, 3, 2),
ConvolutionalLayer(32, 64, 3, 2),
nn.ReLU(),
ConvLayer(64, 128, 3, 2),
ConvolutionalLayer(64, 128, 3, 2),
nn.ReLU()
)
self.ResidualBlock = nn.Sequential(
@@ -136,11 +108,11 @@ class TransformerNetwork(nn.Module):
ResidualLayer(128, 3)
)
self.DeconvBlock = nn.Sequential(
DeconvLayer(128, 64, 3, 2, 1),
DeconvolutionalLayer(128, 64, 3, 2, 1),
nn.ReLU(),
DeconvLayer(64, 32, 3, 2, 1),
DeconvolutionalLayer(64, 32, 3, 2, 1),
nn.ReLU(),
ConvLayer(32, 3, 9, 1, norm="None")
ConvolutionalLayer(32, 3, 9, 1, norm="None")
)
def forward(self, x):
@@ -150,25 +122,27 @@ class TransformerNetwork(nn.Module):
return out
# Node classes
class TrainFastStyleTransfer:
def __init__(self):
pass
@classmethod
def INPUT_TYPES(s):
input_dir = folder_paths.get_input_directory()
files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))]
return {
"required": {
"style_img": (sorted(files), {"image_upload": True}),
"style_img": ("IMAGE",),
"seed": ("INT", {"default": 30, "min": 0, "max": 999999, "step": 1,}),
"content_weight": ("INT", {"default": 14, "min": 1, "max": 128, "step": 1,}),
"style_weight": ("INT", {"default": 50, "min": 1, "max": 128, "step": 1,}),
"batch_size": ("INT", {"default": 4, "min": 1, "max": 128, "step": 1,}),
"train_img_size": ("INT", {"default": 256, "min": 256, "max": 2048, "step": 1,}),
"learning_rate": ("FLOAT", {"default": 0.001, "min": 0.0001, "max": 0.1, "step": 0.0001}),
"tv_weight": ("FLOAT", {"default": 0.001, "min": 0.0, "max": 1.0, "step": 0.0000001}),
"batch_size": ("INT", {"default": 4, "min": 1, "max": 32, "step": 1,}),
"train_img_size": ("INT", {"default": 256, "min": 128, "max": 2048, "step": 1,}),
"learning_rate": ("FLOAT", {"default": 0.001, "min": 0.0001, "max": 100.0, "step": 0.0001}),
"num_epochs": ("INT", {"default": 1, "min": 1, "max": 20, "step": 1,}),
"save_model_every": ("INT", {"default": 500, "min": 100, "max": 10000, "step": 1,}),
"save_model_every": ("INT", {"default": 500, "min": 10, "max": 10000, "step": 10,}),
"from_pretrained": ("INT", {"default": 0, "min": 0, "max": 1, "step": 1,}),
"model": ([file for file in os.listdir(os.path.join(folder_paths.base_path, "custom_nodes/ComfyUI-Fast-Style-Transfer/models/")) if file.endswith('.pth')], ),
},
}
@@ -179,14 +153,19 @@ class TrainFastStyleTransfer:
CATEGORY = "Style Transfer"
def encode_tensor(self, tensor):
tensor = tensor.permute(0, 3, 1, 2).contiguous() # Convert to [batch_size, channels, height, width]
return tensor[:, [2, 1, 0], :, :] * 255
def train(self, style_img, seed, batch_size, train_img_size, learning_rate, num_epochs, content_weight, style_weight, save_model_every):
def train(self, style_img, seed, batch_size, train_img_size, learning_rate, num_epochs, content_weight, style_weight, tv_weight, save_model_every, from_pretrained, model):
temp_save_style_img = os.path.join(folder_paths.base_path, "custom_nodes/ComfyUI-Fast-Style-Transfer/temp/") + "temp_save_content_img.pt"
save_model_path = os.path.join(folder_paths.base_path, "custom_nodes/ComfyUI-Fast-Style-Transfer/models/")
dataset_path = os.path.join(folder_paths.base_path, "custom_nodes/ComfyUI-Fast-Style-Transfer/dataset/")
vgg_path = os.path.join(folder_paths.base_path, "custom_nodes/ComfyUI-Fast-Style-Transfer/vgg/vgg16-00b39a1b.pth")
save_image_path = os.path.join(folder_paths.base_path, "custom_nodes/ComfyUI-Fast-Style-Transfer/output/")
style_image_path = folder_paths.get_annotated_filepath(style_img)
train_path = os.path.join(folder_paths.base_path, "custom_nodes/ComfyUI-Fast-Style-Transfer/train.py")
command = [
@@ -195,24 +174,31 @@ class TrainFastStyleTransfer:
'--dataset_path', dataset_path,
'--vgg_path', vgg_path,
'--num_epochs', str(num_epochs),
'--style_image_path', style_image_path,
'--temp_save_style_img', temp_save_style_img,
'--batch_size', str(batch_size),
'--content_weight', str(content_weight),
'--style_weight', str(style_weight),
'--tv_weight', str(tv_weight),
'--adam_lr', str(learning_rate),
'--save_model_path', save_model_path,
'--save_image_path', save_image_path,
'--save_model_every', str(save_model_every),
'--seed', str(seed)
'--seed', str(seed),
'--pretrained_model'
]
if from_pretrained:
command.append(model)
else:
command.append('none')
torch.save(self.encode_tensor(style_img), temp_save_style_img)
sp.run(command)
return ()
class FastStyleTransfer:
def __init__(self):
pass
@@ -262,15 +248,81 @@ class FastStyleTransfer:
return (image,)
class NeuralStyleTransfer:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"content_img": ("IMAGE",),
"style_img": ("IMAGE",),
"content_weight": ("FLOAT", {"default": 1e5, "min": 1e3, "max": 1e6, "step": 1e3}),
"style_weight": ("FLOAT", {"default": 3e4, "min": 1e1, "max": 1e5, "step": 1e1}),
"tv_weight": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1e1, "step": 0.1}),
"num_steps": ("INT", {"default": 100, "min": 10, "max": 10000, "step": 10}),
"learning_rate": ("FLOAT", {"default": 1.0, "min": 1e-4, "max": 1e3, "step": 0.1}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "neural_style_transfer"
CATEGORY = "Style Transfer"
def encode_tensor(self, tensor):
tensor = tensor.permute(0, 3, 1, 2).contiguous() # Convert to [batch_size, channels, height, width]
return tensor * 255
def decode_tensor(self, tensor):
tensor = tensor.permute(0, 2, 3, 1).contiguous() # Convert to [batch_size, height, width, channels]
return tensor / 255
def neural_style_transfer(self, content_img, style_img, content_weight, style_weight, tv_weight, num_steps, learning_rate):
neural_style_transfer_path = os.path.join(folder_paths.base_path, "custom_nodes/ComfyUI-Fast-Style-Transfer/neural_style_transfer.py")
temp_save_content_img = os.path.join(folder_paths.base_path, "custom_nodes/ComfyUI-Fast-Style-Transfer/temp/") + "temp_save_content_img.pt"
temp_save_style_img = os.path.join(folder_paths.base_path, "custom_nodes/ComfyUI-Fast-Style-Transfer/temp/") + "temp_save_style_img.pt"
temp_load_final_img = os.path.join(folder_paths.base_path, "custom_nodes/ComfyUI-Fast-Style-Transfer/temp/") + "temp_load_final_img.pt"
torch.save(self.encode_tensor(content_img), temp_save_content_img)
torch.save(self.encode_tensor(style_img), temp_save_style_img)
command = [
sys.executable, neural_style_transfer_path,
'--content_weight', str(content_weight),
'--style_weight', str(style_weight),
'--tv_weight', str(tv_weight),
'--temp_save_style_img', temp_save_style_img,
'--temp_save_content_img', temp_save_content_img,
'--temp_load_final_img', temp_load_final_img,
'--num_steps', str(num_steps),
'--learning_rate', str(learning_rate)
]
sp.run(command)
image = self.decode_tensor(torch.load(temp_load_final_img))
os.remove(temp_save_style_img)
os.remove(temp_save_content_img)
os.remove(temp_load_final_img)
return (image,)
# A dictionary that contains all nodes you want to export with their names
# NOTE: names should be globally unique
NODE_CLASS_MAPPINGS = {
"FastStyleTransfer": FastStyleTransfer,
"TrainFastStyleTransfer": TrainFastStyleTransfer
"TrainFastStyleTransfer": TrainFastStyleTransfer,
"NeuralStyleTransfer": NeuralStyleTransfer,
}
# A dictionary that contains the friendly/humanly readable titles for the nodes
NODE_DISPLAY_NAME_MAPPINGS = {
"FastStyleTransfer": "Fast Style Transfer",
"TrainFastStyleTransfer": "Train Fast Style Transfer"
"TrainFastStyleTransfer": "Train Fast Style Transfer",
"NeuralStyleTransfer": "Neural Style Transfer",
}