Add files via upload
This commit is contained in:
zeroxoxo
+111
-59
@@ -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",
|
||||
}
|
||||
|
||||
@@ -0,0 +1,167 @@
|
||||
import torch
|
||||
from torch.optim import AdamW, LBFGS
|
||||
from torchvision import models
|
||||
import torch.nn as nn
|
||||
import argparse
|
||||
from collections import namedtuple
|
||||
|
||||
class Vgg16(torch.nn.Module):
|
||||
"""Only those layers are exposed which have already proven to work nicely."""
|
||||
def __init__(self, requires_grad=False, show_progress=False):
|
||||
super().__init__()
|
||||
vgg_pretrained_features = models.vgg16(pretrained=True, progress=show_progress).features
|
||||
self.layer_names = ['relu1_2', 'relu2_2', 'relu3_3', 'relu4_3']
|
||||
self.content_feature_maps_index = 1 # relu2_2
|
||||
self.style_feature_maps_indices = list(range(len(self.layer_names))) # all layers used for style representation
|
||||
|
||||
self.slice1 = torch.nn.Sequential()
|
||||
self.slice2 = torch.nn.Sequential()
|
||||
self.slice3 = torch.nn.Sequential()
|
||||
self.slice4 = torch.nn.Sequential()
|
||||
for x in range(4):
|
||||
self.slice1.add_module(str(x), vgg_pretrained_features[x])
|
||||
for x in range(4, 9):
|
||||
self.slice2.add_module(str(x), vgg_pretrained_features[x])
|
||||
for x in range(9, 16):
|
||||
self.slice3.add_module(str(x), vgg_pretrained_features[x])
|
||||
for x in range(16, 23):
|
||||
self.slice4.add_module(str(x), vgg_pretrained_features[x])
|
||||
if not requires_grad:
|
||||
for param in self.parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
def forward(self, x):
|
||||
x = self.slice1(x)
|
||||
relu1_2 = x
|
||||
x = self.slice2(x)
|
||||
relu2_2 = x
|
||||
x = self.slice3(x)
|
||||
relu3_3 = x
|
||||
x = self.slice4(x)
|
||||
relu4_3 = x
|
||||
vgg_outputs = namedtuple("VggOutputs", self.layer_names)
|
||||
out = vgg_outputs(relu1_2, relu2_2, relu3_3, relu4_3)
|
||||
return out
|
||||
|
||||
|
||||
def prepare_model(device):
|
||||
model = Vgg16()
|
||||
|
||||
content_feature_maps_index = model.content_feature_maps_index
|
||||
style_feature_maps_indices = model.style_feature_maps_indices
|
||||
layer_names = model.layer_names
|
||||
|
||||
content_fms_index_name = (content_feature_maps_index, layer_names[content_feature_maps_index])
|
||||
style_fms_indices_names = (style_feature_maps_indices, layer_names)
|
||||
return model.to(device).eval(), content_fms_index_name, style_fms_indices_names
|
||||
|
||||
|
||||
def gram_matrix(x, should_normalize=True):
|
||||
(b, ch, h, w) = x.size()
|
||||
features = x.view(b, ch, w * h)
|
||||
features_t = features.transpose(1, 2)
|
||||
gram = features.bmm(features_t)
|
||||
if should_normalize:
|
||||
gram /= ch * h * w
|
||||
return gram
|
||||
|
||||
|
||||
def total_variation(y):
|
||||
return torch.sum(torch.abs(y[:, :, :, :-1] - y[:, :, :, 1:])) + \
|
||||
torch.sum(torch.abs(y[:, :, :-1, :] - y[:, :, 1:, :]))
|
||||
|
||||
|
||||
def build_loss(neural_net, optimizing_img, target_representations, content_feature_maps_index, style_feature_maps_indices, config):
|
||||
target_content_representation = target_representations[0]
|
||||
target_style_representation = target_representations[1]
|
||||
|
||||
current_set_of_feature_maps = neural_net(optimizing_img)
|
||||
|
||||
current_content_representation = current_set_of_feature_maps[content_feature_maps_index].squeeze(axis=0)
|
||||
content_loss = torch.nn.MSELoss(reduction='mean')(target_content_representation, current_content_representation)
|
||||
|
||||
style_loss = 0.0
|
||||
current_style_representation = [gram_matrix(x) for cnt, x in enumerate(current_set_of_feature_maps) if cnt in style_feature_maps_indices]
|
||||
for gram_gt, gram_hat in zip(target_style_representation, current_style_representation):
|
||||
style_loss += torch.nn.MSELoss(reduction='sum')(gram_gt[0], gram_hat[0])
|
||||
style_loss /= len(target_style_representation)
|
||||
|
||||
tv_loss = total_variation(optimizing_img)
|
||||
|
||||
total_loss = config['content_weight'] * content_loss + config['style_weight'] * style_loss + config['tv_weight'] * tv_loss
|
||||
|
||||
return total_loss, content_loss, style_loss, tv_loss
|
||||
|
||||
def make_tuning_step(neural_net, optimizer, target_representations, content_feature_maps_index, style_feature_maps_indices, config):
|
||||
def tuning_step(optimizing_img):
|
||||
optimizer.zero_grad()
|
||||
total_loss, content_loss, style_loss, tv_loss = build_loss(neural_net, optimizing_img, target_representations, content_feature_maps_index, style_feature_maps_indices, config)
|
||||
total_loss.backward()
|
||||
optimizer.step()
|
||||
return total_loss, content_loss, style_loss, tv_loss
|
||||
return tuning_step
|
||||
|
||||
def neural_style_transfer_from_tensors(content_img_tensor, style_img_tensor, config, num_steps, learning_rate):
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
content_img = content_img_tensor.to(device)
|
||||
style_img = style_img_tensor.to(device)
|
||||
|
||||
optimizing_img = content_img.clone().requires_grad_(True)
|
||||
|
||||
neural_net, content_feature_maps_index_name, style_feature_maps_indices_names = prepare_model(device)
|
||||
|
||||
content_img_set_of_feature_maps = neural_net(content_img)
|
||||
style_img_set_of_feature_maps = neural_net(style_img)
|
||||
|
||||
target_content_representation = content_img_set_of_feature_maps[content_feature_maps_index_name[0]].squeeze(axis=0)
|
||||
target_style_representation = [gram_matrix(x) for cnt, x in enumerate(style_img_set_of_feature_maps) if cnt in style_feature_maps_indices_names[0]]
|
||||
target_representations = [target_content_representation, target_style_representation]
|
||||
|
||||
optimizer = LBFGS((optimizing_img,), max_iter=num_steps, lr=learning_rate, line_search_fn='strong_wolfe')
|
||||
cnt = 0
|
||||
|
||||
def closure():
|
||||
nonlocal cnt
|
||||
if torch.is_grad_enabled():
|
||||
optimizer.zero_grad()
|
||||
total_loss, content_loss, style_loss, tv_loss = build_loss(neural_net, optimizing_img, target_representations, content_feature_maps_index_name[0], style_feature_maps_indices_names[0], config)
|
||||
if total_loss.requires_grad:
|
||||
total_loss.backward()
|
||||
if cnt%100==0:
|
||||
with torch.no_grad():
|
||||
print(f'L-BFGS | iteration: {cnt:03}, total loss={total_loss.item():12.4f}, content_loss={config["content_weight"] * content_loss.item():12.4f}, style loss={config["style_weight"] * style_loss.item():12.4f}, tv loss={config["tv_weight"] * tv_loss.item():12.4f}')
|
||||
|
||||
cnt += 1
|
||||
return total_loss
|
||||
|
||||
optimizer.step(closure)
|
||||
|
||||
return optimizing_img.detach()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--temp_save_style_img", type=str)
|
||||
parser.add_argument("--temp_save_content_img", type=str)
|
||||
parser.add_argument("--temp_load_final_img", type=str)
|
||||
|
||||
parser.add_argument("--content_weight", type=float, help="weight factor for content loss", default=1e5)
|
||||
parser.add_argument("--style_weight", type=float, help="weight factor for style loss", default=3e4)
|
||||
parser.add_argument("--tv_weight", type=float, help="weight factor for total variation loss", default=1e0)
|
||||
parser.add_argument("--learning_rate", type=float, help="learning_rate", default=1e1)
|
||||
parser.add_argument("--num_steps", type=int, help="number of training steps", default=100)
|
||||
|
||||
args = parser.parse_args()
|
||||
config = {
|
||||
"content_weight": args.content_weight,
|
||||
"style_weight": args.style_weight,
|
||||
"tv_weight": args.tv_weight,
|
||||
}
|
||||
style_img = torch.load(args.temp_save_style_img)
|
||||
content_img = torch.load(args.temp_save_content_img)
|
||||
final_img = neural_style_transfer_from_tensors(content_img, style_img, config, args.num_steps, args.learning_rate)
|
||||
print("nst done")
|
||||
|
||||
torch.save(final_img, args.temp_load_final_img)
|
||||
|
||||
@@ -7,8 +7,6 @@ import numpy as np
|
||||
import time
|
||||
import argparse
|
||||
import cv2
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
|
||||
parser = argparse.ArgumentParser(description="Train a neural network with style transfer.")
|
||||
|
||||
@@ -16,15 +14,17 @@ parser.add_argument('--train_image_size', type=int, default=256, help='Size of t
|
||||
parser.add_argument('--dataset_path', type=str, default="dataset", help='Path to the dataset')
|
||||
parser.add_argument('--vgg_path', type=str, default="vgg", help='Path to the vgg model')
|
||||
parser.add_argument('--num_epochs', type=int, default=1, help='Number of epochs for training')
|
||||
parser.add_argument('--style_image_path', type=str, default="images/oil.jpg", help='Path to the style image')
|
||||
parser.add_argument('--temp_save_style_img', type=str, default="temp/temp_save_style_img.pt", help='Path to the style image tensor')
|
||||
parser.add_argument('--batch_size', type=int, default=12, help='Batch size for training')
|
||||
parser.add_argument('--content_weight', type=float, default=8, help='Weight for content loss')
|
||||
parser.add_argument('--style_weight', type=float, default=50, help='Weight for style loss')
|
||||
parser.add_argument('--tv_weight', type=float, default=0.001, help='Weight for total variation loss')
|
||||
parser.add_argument('--adam_lr', type=float, default=0.001, help='Learning rate for Adam optimizer')
|
||||
parser.add_argument('--save_model_path', type=str, default="models/oil/", help='Path to save the trained model')
|
||||
parser.add_argument('--save_image_path', type=str, default="images/out/", help='Path to save the output images')
|
||||
parser.add_argument('--save_model_every', type=int, default=200, help='Save model every n batches')
|
||||
parser.add_argument('--seed', type=int, default=1234, help='Random seed')
|
||||
parser.add_argument('--pretrained_model', type=str, default='none', help='pretrained model')
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
@@ -32,89 +32,34 @@ args = parser.parse_args()
|
||||
TRAIN_IMAGE_SIZE = args.train_image_size
|
||||
DATASET_PATH = args.dataset_path
|
||||
NUM_EPOCHS = args.num_epochs
|
||||
STYLE_IMAGE_PATH = args.style_image_path
|
||||
STYLE_IMAGE_PATH = args.temp_save_style_img
|
||||
VGG_PATH = args.vgg_path
|
||||
BATCH_SIZE = args.batch_size
|
||||
CONTENT_WEIGHT = args.content_weight
|
||||
STYLE_WEIGHT = args.style_weight
|
||||
TV_WEIGHT = args.tv_weight
|
||||
ADAM_LR = args.adam_lr
|
||||
SAVE_MODEL_PATH = args.save_model_path
|
||||
SAVE_IMAGE_PATH = args.save_image_path
|
||||
SAVE_MODEL_EVERY = args.save_model_every
|
||||
SEED = args.seed
|
||||
PLOT_LOSS = 1
|
||||
MODEL = args.pretrained_model
|
||||
|
||||
# Utils
|
||||
|
||||
def plot_loss_hist(c_loss, s_loss, total_loss, title="Loss History"):
|
||||
x = [i for i in range(len(total_loss))]
|
||||
plt.figure(figsize=[10, 6])
|
||||
plt.plot(x, c_loss, label="Content Loss")
|
||||
plt.plot(x, s_loss, label="Style Loss")
|
||||
plt.plot(x, total_loss, label="Total Loss")
|
||||
|
||||
plt.legend()
|
||||
plt.xlabel('Every 500 iterations')
|
||||
plt.ylabel('Loss')
|
||||
plt.title(title)
|
||||
plt.show()
|
||||
|
||||
# Gram Matrix
|
||||
def gram(tensor):
|
||||
B, C, H, W = tensor.shape
|
||||
x = tensor.view(B, C, H*W)
|
||||
x_t = x.transpose(1, 2)
|
||||
return torch.bmm(x, x_t) / (C*H*W)
|
||||
|
||||
# Load image file
|
||||
def load_image(path):
|
||||
# Images loaded as BGR
|
||||
img = cv2.imread(path)
|
||||
return img
|
||||
|
||||
# Show image
|
||||
def show(img):
|
||||
# Convert from BGR to RGB
|
||||
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
||||
|
||||
# imshow() only accepts float [0,1] or int [0,255]
|
||||
img = np.array(img/255).clip(0,1)
|
||||
|
||||
plt.figure(figsize=(10, 5))
|
||||
plt.imshow(img)
|
||||
plt.show()
|
||||
|
||||
gram_matrix = torch.bmm(x, x_t) / (C*H*W)
|
||||
gram_matrix = torch.clamp(gram_matrix, min=1e-6, max=1e6)
|
||||
return gram_matrix
|
||||
|
||||
# Save image
|
||||
def saveimg(img, image_path):
|
||||
img = img.clip(0, 255)
|
||||
cv2.imwrite(image_path, img)
|
||||
|
||||
# Preprocessing ~ Image to Tensor
|
||||
def itot(img, max_size=None):
|
||||
# Rescale the image
|
||||
if (max_size==None):
|
||||
itot_t = transforms.Compose([
|
||||
#transforms.ToPILImage(),
|
||||
transforms.ToTensor(),
|
||||
transforms.Lambda(lambda x: x.mul(255))
|
||||
])
|
||||
else:
|
||||
H, W, C = img.shape
|
||||
image_size = tuple([int((float(max_size) / max([H,W]))*x) for x in [H, W]])
|
||||
itot_t = transforms.Compose([
|
||||
transforms.ToPILImage(),
|
||||
transforms.Resize(image_size),
|
||||
transforms.ToTensor(),
|
||||
transforms.Lambda(lambda x: x.mul(255))
|
||||
])
|
||||
|
||||
# Convert image to tensor
|
||||
tensor = itot_t(img)
|
||||
|
||||
# Add the batch_size dimension
|
||||
tensor = tensor.unsqueeze(dim=0)
|
||||
return tensor
|
||||
|
||||
# Preprocessing ~ Tensor to Image
|
||||
def ttoi(tensor):
|
||||
# Add the means
|
||||
@@ -130,6 +75,11 @@ def ttoi(tensor):
|
||||
img = img.transpose(1, 2, 0)
|
||||
return img
|
||||
|
||||
# Alternative loss functions experiments
|
||||
def total_variation(y):
|
||||
return torch.sum(torch.abs(y[:, :, :, :-1] - y[:, :, :, 1:])) + \
|
||||
torch.sum(torch.abs(y[:, :, :-1, :] - y[:, :, 1:, :]))
|
||||
|
||||
# VGG
|
||||
class VGG16(nn.Module):
|
||||
def __init__(self, vgg_path="models/vgg16-00b39a1b.pth"):
|
||||
@@ -164,11 +114,11 @@ class TransformerNetworkClass(nn.Module):
|
||||
def __init__(self):
|
||||
super(TransformerNetworkClass, 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(
|
||||
@@ -179,11 +129,11 @@ class TransformerNetworkClass(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):
|
||||
@@ -192,9 +142,9 @@ class TransformerNetworkClass(nn.Module):
|
||||
out = self.DeconvBlock(x)
|
||||
return out
|
||||
|
||||
class ConvLayer(nn.Module):
|
||||
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
|
||||
padding_size = kernel_size // 2
|
||||
self.reflection_pad = nn.ReflectionPad2d(padding_size)
|
||||
@@ -204,6 +154,7 @@ class ConvLayer(nn.Module):
|
||||
|
||||
# Normalization Layers
|
||||
self.norm_type = norm
|
||||
|
||||
if (norm=="instance"):
|
||||
self.norm_layer = nn.InstanceNorm2d(out_channels, affine=True)
|
||||
elif (norm=="batch"):
|
||||
@@ -226,9 +177,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
|
||||
@@ -237,9 +188,9 @@ class ResidualLayer(nn.Module):
|
||||
out = out + identity # add residual
|
||||
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
|
||||
@@ -247,6 +198,7 @@ class DeconvLayer(nn.Module):
|
||||
|
||||
# Normalization Layers
|
||||
self.norm_type = norm
|
||||
|
||||
if (norm=="instance"):
|
||||
self.norm_layer = nn.InstanceNorm2d(out_channels, affine=True)
|
||||
elif (norm=="batch"):
|
||||
@@ -261,8 +213,8 @@ class DeconvLayer(nn.Module):
|
||||
return out
|
||||
|
||||
|
||||
|
||||
def train():
|
||||
|
||||
# Seeds
|
||||
torch.manual_seed(SEED)
|
||||
torch.cuda.manual_seed(SEED)
|
||||
@@ -286,15 +238,16 @@ def train():
|
||||
|
||||
# Load networks
|
||||
TransformerNetwork = TransformerNetworkClass().to(device)
|
||||
if MODEL != 'none':
|
||||
model_path = SAVE_MODEL_PATH + MODEL
|
||||
print(f"Loading {model_path} model")
|
||||
TransformerNetwork.load_state_dict(torch.load(model_path, map_location=device))
|
||||
TransformerNetwork.to(device)
|
||||
VGG = VGG16(vgg_path=VGG_PATH).to(device)
|
||||
|
||||
# Get Style Features
|
||||
imagenet_neg_mean = torch.tensor([-103.939, -116.779, -123.68], dtype=torch.float32).reshape(1,3,1,1).to(device)
|
||||
print(f"Train.py: style image path: {STYLE_IMAGE_PATH}")
|
||||
style_image = load_image(STYLE_IMAGE_PATH)
|
||||
print(f"Train.py: style image type: {type(style_image)}")
|
||||
style_tensor = itot(style_image).to(device)
|
||||
print(f"Train.py: style tensor type: {type(style_tensor)}")
|
||||
style_tensor = torch.load(STYLE_IMAGE_PATH).to(device)
|
||||
style_tensor = style_tensor.add(imagenet_neg_mean)
|
||||
B, C, H, W = style_tensor.shape
|
||||
style_features = VGG(style_tensor.expand([BATCH_SIZE, C, H, W]))
|
||||
@@ -304,7 +257,7 @@ def train():
|
||||
|
||||
# Optimizer settings
|
||||
optimizer = optim.AdamW(TransformerNetwork.parameters(), lr=ADAM_LR, fused=True, amsgrad=True)
|
||||
|
||||
scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=20, T_mult=2, eta_min=1e-6, last_epoch=-1)
|
||||
# Loss trackers
|
||||
content_loss_history = []
|
||||
style_loss_history = []
|
||||
@@ -312,22 +265,13 @@ def train():
|
||||
batch_content_loss_sum = 0
|
||||
batch_style_loss_sum = 0
|
||||
batch_total_loss_sum = 0
|
||||
|
||||
# Optimization/Training Loop
|
||||
batch_count = 1
|
||||
start_time = time.time()
|
||||
for epoch in range(NUM_EPOCHS):
|
||||
print("========Epoch {}/{}========".format(epoch+1, NUM_EPOCHS))
|
||||
for content_batch, _ in train_loader:
|
||||
# Get current batch size in case of odd batch sizes
|
||||
curr_batch_size = content_batch.shape[0]
|
||||
|
||||
# Free-up unneeded cuda memory
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
# Zero-out Gradients
|
||||
for batch_count, (content_batch, _) in enumerate(train_loader):
|
||||
optimizer.zero_grad()
|
||||
|
||||
torch.cuda.empty_cache()
|
||||
# Generate images and get features
|
||||
content_batch = content_batch[:,[2,1,0]].to(device)
|
||||
generated_batch = TransformerNetwork(content_batch)
|
||||
@@ -336,55 +280,63 @@ def train():
|
||||
|
||||
# Content Loss
|
||||
MSELoss = nn.MSELoss().to(device)
|
||||
content_loss = CONTENT_WEIGHT * MSELoss(generated_features['relu2_2'], content_features['relu2_2'])
|
||||
MAELoss = nn.L1Loss().to(device)
|
||||
content_loss = CONTENT_WEIGHT * MSELoss(generated_features['relu2_2'], content_features['relu2_2']) + MAELoss(generated_features['relu2_2'], content_features['relu2_2']) / 2
|
||||
batch_content_loss_sum += content_loss
|
||||
|
||||
# Style Loss
|
||||
style_loss = 0
|
||||
style_loss = 0.0
|
||||
for key, value in generated_features.items():
|
||||
s_loss = MSELoss(gram(value), style_gram[key][:curr_batch_size])
|
||||
style_loss += s_loss
|
||||
style_loss *= STYLE_WEIGHT
|
||||
style_loss += MSELoss(gram(value), style_gram[key][:content_batch.shape[0]])
|
||||
style_loss += MAELoss(gram(value), style_gram[key][:content_batch.shape[0]])
|
||||
style_loss *= STYLE_WEIGHT/2
|
||||
batch_style_loss_sum += style_loss.item()
|
||||
|
||||
# Total Loss
|
||||
total_loss = content_loss + style_loss
|
||||
batch_total_loss_sum += total_loss.item()
|
||||
# TV loss
|
||||
tv_loss = total_variation(generated_batch) * TV_WEIGHT
|
||||
|
||||
# Backprop and Weight Update
|
||||
# Total Loss
|
||||
total_loss = content_loss + style_loss + tv_loss
|
||||
batch_total_loss_sum += total_loss.item()
|
||||
|
||||
total_loss.backward()
|
||||
optimizer.step()
|
||||
if batch_count % 50 == 0:
|
||||
scheduler.step()
|
||||
|
||||
with torch.no_grad():
|
||||
print(f'AdamW | iteration: {batch_count+1:03}, total loss={total_loss.item():12.4f}, content_loss={content_loss.item():12.4f}, style loss={style_loss.item():12.4f}, tv loss={tv_loss.item():12.4f}')
|
||||
|
||||
# Save Model and Print Losses
|
||||
if (((batch_count-1)%SAVE_MODEL_EVERY == 0) or (batch_count==NUM_EPOCHS*len(train_loader))):
|
||||
if (((batch_count)%SAVE_MODEL_EVERY == 0) or ((batch_count+1)==NUM_EPOCHS*len(train_loader))):
|
||||
# Print Losses
|
||||
print("========Iteration {}/{}========".format(batch_count, NUM_EPOCHS*len(train_loader)))
|
||||
print("\tContent Loss:\t{:.2f}".format(batch_content_loss_sum/batch_count))
|
||||
print("\tStyle Loss:\t{:.2f}".format(batch_style_loss_sum/batch_count))
|
||||
print("\tTotal Loss:\t{:.2f}".format(batch_total_loss_sum/batch_count))
|
||||
print("========Iteration {}/{}========".format(batch_count+1, NUM_EPOCHS*len(train_loader)))
|
||||
print("\tContent Loss:\t{:.2f}".format(batch_content_loss_sum/(batch_count+1)))
|
||||
print("\tStyle Loss:\t{:.2f}".format(batch_style_loss_sum/(batch_count+1)))
|
||||
print("\tTotal Loss:\t{:.2f}".format(batch_total_loss_sum/(batch_count+1)))
|
||||
print("Time elapsed:\t{} seconds".format(time.time()-start_time))
|
||||
|
||||
# Save Model
|
||||
checkpoint_path = SAVE_MODEL_PATH + "checkpoint_" + str(batch_count-1) + ".pth"
|
||||
checkpoint_path = SAVE_MODEL_PATH + "checkpoint_" + str((batch_count+1)) + ".pth"
|
||||
torch.save(TransformerNetwork.state_dict(), checkpoint_path)
|
||||
print("Saved TransformerNetwork checkpoint file at {}".format(checkpoint_path))
|
||||
|
||||
# Save sample generated image
|
||||
sample_tensor = generated_batch[0].clone().detach().unsqueeze(dim=0)
|
||||
sample_image = ttoi(sample_tensor.clone().detach())
|
||||
sample_image_path = SAVE_IMAGE_PATH + "sample0_" + str(batch_count-1) + ".png"
|
||||
sample_image_path = SAVE_IMAGE_PATH + "sample0_" + str((batch_count+1)) + ".png"
|
||||
saveimg(sample_image, sample_image_path)
|
||||
print("Saved sample tranformed image at {}".format(sample_image_path))
|
||||
|
||||
# Save loss histories
|
||||
content_loss_history.append(batch_content_loss_sum/batch_count)
|
||||
style_loss_history.append(batch_style_loss_sum/batch_count)
|
||||
total_loss_history.append(batch_total_loss_sum/batch_count)
|
||||
content_loss_history.append(batch_content_loss_sum/(batch_count+1))
|
||||
style_loss_history.append(batch_style_loss_sum/(batch_count+1))
|
||||
total_loss_history.append(batch_total_loss_sum/(batch_count+1))
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
# Iterate Batch Counter
|
||||
batch_count+=1
|
||||
if batch_count%10 == 0:
|
||||
print(f"Batch count is {batch_count}")
|
||||
|
||||
stop_time = time.time()
|
||||
# Print loss histories
|
||||
@@ -397,17 +349,7 @@ def train():
|
||||
print("========Total Loss========")
|
||||
print(total_loss_history)
|
||||
|
||||
# Save TransformerNetwork weights
|
||||
TransformerNetwork.eval()
|
||||
TransformerNetwork.cpu()
|
||||
final_path = SAVE_MODEL_PATH + "transformer_weight.pth"
|
||||
print("Saving TransformerNetwork weights at {}".format(final_path))
|
||||
torch.save(TransformerNetwork.state_dict(), final_path)
|
||||
print("Done saving final model")
|
||||
|
||||
# Plot Loss Histories
|
||||
if (PLOT_LOSS):
|
||||
plot_loss_hist(content_loss_history, style_loss_history, total_loss_history)
|
||||
|
||||
if __name__ == "__main__":
|
||||
train()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user