From 968e445a9ee96187877f9a896720f7a5fdc01e1a Mon Sep 17 00:00:00 2001 From: zeroxoxo <32829385+zeroxoxo@users.noreply.github.com> Date: Sun, 30 Jun 2024 18:41:04 +0500 Subject: [PATCH] Add TrainFastStyleTransfer node --- FastStyleTransferNode.py | 106 +++++++++- train.py | 413 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 513 insertions(+), 6 deletions(-) create mode 100644 train.py diff --git a/FastStyleTransferNode.py b/FastStyleTransferNode.py index 35fd3f4..77a2766 100644 --- a/FastStyleTransferNode.py +++ b/FastStyleTransferNode.py @@ -5,19 +5,49 @@ https://github.com/rrmina/fast-neural-style-pytorch.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): def __init__(self, in_channels, out_channels, kernel_size, stride, norm="instance"): super(ConvLayer, self).__init__() # Padding Layers - padding_size = kernel_size // 2 - self.reflection_pad = nn.ReflectionPad2d(padding_size) + self.padding_size = kernel_size // 2 + self.reflection_pad = nn.ReflectionPad2d(self.padding_size) # Convolution Layer self.conv_layer = nn.Conv2d(in_channels, out_channels, kernel_size, stride) @@ -32,7 +62,7 @@ class ConvLayer(nn.Module): def forward(self, x): x = self.reflection_pad(x) x = self.conv_layer(x) - if (self.norm_type=="None"): + if self.norm_type == "None": out = x else: out = self.norm_layer(x) @@ -120,6 +150,68 @@ class TransformerNetwork(nn.Module): return out +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}), + "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}), + "num_epochs": ("INT", {"default": 1, "min": 1, "max": 20, "step": 1,}), + "save_model_every": ("INT", {"default": 500, "min": 100, "max": 10000, "step": 1,}), + }, + } + + RETURN_TYPES = () + OUTPUT_NODE = True + + FUNCTION = "train" + + CATEGORY = "Style Transfer" + + + def train(self, style_img, seed, batch_size, train_img_size, learning_rate, num_epochs, content_weight, style_weight, save_model_every): + 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 = [ + sys.executable, train_path, + '--train_image_size', str(train_img_size), + '--dataset_path', dataset_path, + '--vgg_path', vgg_path, + '--num_epochs', str(num_epochs), + '--style_image_path', style_image_path, + '--batch_size', str(batch_size), + '--content_weight', str(content_weight), + '--style_weight', str(style_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) + ] + + sp.run(command) + return () + + + + class FastStyleTransfer: def __init__(self): @@ -173,10 +265,12 @@ class FastStyleTransfer: # A dictionary that contains all nodes you want to export with their names # NOTE: names should be globally unique NODE_CLASS_MAPPINGS = { - "FastStyleTransfer": FastStyleTransfer + "FastStyleTransfer": FastStyleTransfer, + "TrainFastStyleTransfer": TrainFastStyleTransfer } # A dictionary that contains the friendly/humanly readable titles for the nodes NODE_DISPLAY_NAME_MAPPINGS = { - "FastStyleTransfer": " Fast Style Transfer" + "FastStyleTransfer": "Fast Style Transfer", + "TrainFastStyleTransfer": "Train Fast Style Transfer" } diff --git a/train.py b/train.py new file mode 100644 index 0000000..a51bb7a --- /dev/null +++ b/train.py @@ -0,0 +1,413 @@ +import torch +import torch.nn as nn +import torch.optim as optim +from torchvision import datasets, transforms, models +import random +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.") + +parser.add_argument('--train_image_size', type=int, default=256, help='Size of the training images') +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('--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('--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') + +args = parser.parse_args() + +# GLOBAL SETTINGS +TRAIN_IMAGE_SIZE = args.train_image_size +DATASET_PATH = args.dataset_path +NUM_EPOCHS = args.num_epochs +STYLE_IMAGE_PATH = args.style_image_path +VGG_PATH = args.vgg_path +BATCH_SIZE = args.batch_size +CONTENT_WEIGHT = args.content_weight +STYLE_WEIGHT = args.style_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 + +# 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() + + +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 + #ttoi_t = transforms.Compose([ + # transforms.Normalize([-103.939, -116.779, -123.68],[1,1,1])]) + + # Remove the batch_size dimension + tensor = tensor.squeeze() + #img = ttoi_t(tensor) + img = tensor.cpu().numpy() + + # Transpose from [C, H, W] -> [H, W, C] + img = img.transpose(1, 2, 0) + return img + +# VGG +class VGG16(nn.Module): + def __init__(self, vgg_path="models/vgg16-00b39a1b.pth"): + 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 + for param in self.features.parameters(): + param.requires_grad = False + + 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 + +# Transformer +class TransformerNetworkClass(nn.Module): + """Feedforward Transformation Network without Tanh + reference: https://arxiv.org/abs/1603.08155 + exact architecture: https://cs.stanford.edu/people/jcjohns/papers/fast-style/fast-style-supp.pdf + """ + def __init__(self): + super(TransformerNetworkClass, self).__init__() + self.ConvBlock = nn.Sequential( + ConvLayer(3, 32, 9, 1), + nn.ReLU(), + ConvLayer(32, 64, 3, 2), + nn.ReLU(), + ConvLayer(64, 128, 3, 2), + nn.ReLU() + ) + self.ResidualBlock = nn.Sequential( + ResidualLayer(128, 3), + ResidualLayer(128, 3), + ResidualLayer(128, 3), + ResidualLayer(128, 3), + ResidualLayer(128, 3) + ) + self.DeconvBlock = nn.Sequential( + DeconvLayer(128, 64, 3, 2, 1), + nn.ReLU(), + DeconvLayer(64, 32, 3, 2, 1), + nn.ReLU(), + ConvLayer(32, 3, 9, 1, norm="None") + ) + + def forward(self, x): + x = self.ConvBlock(x) + x = self.ResidualBlock(x) + out = self.DeconvBlock(x) + return out + +class ConvLayer(nn.Module): + def __init__(self, in_channels, out_channels, kernel_size, stride, norm="instance"): + super(ConvLayer, self).__init__() + # Padding Layers + padding_size = kernel_size // 2 + self.reflection_pad = nn.ReflectionPad2d(padding_size) + + # Convolution Layer + self.conv_layer = nn.Conv2d(in_channels, out_channels, kernel_size, stride) + + # Normalization Layers + self.norm_type = norm + if (norm=="instance"): + self.norm_layer = nn.InstanceNorm2d(out_channels, affine=True) + elif (norm=="batch"): + self.norm_layer = nn.BatchNorm2d(out_channels, affine=True) + + def forward(self, x): + x = self.reflection_pad(x) + x = self.conv_layer(x) + if (self.norm_type=="None"): + out = x + else: + out = self.norm_layer(x) + return out + +class ResidualLayer(nn.Module): + """ + Deep Residual Learning for Image Recognition + + https://arxiv.org/abs/1512.03385 + """ + def __init__(self, channels=128, kernel_size=3): + super(ResidualLayer, self).__init__() + self.conv1 = ConvLayer(channels, channels, kernel_size, stride=1) + self.relu = nn.ReLU() + self.conv2 = ConvLayer(channels, channels, kernel_size, stride=1) + + def forward(self, x): + identity = x # preserve residual + out = self.relu(self.conv1(x)) # 1st conv layer + activation + out = self.conv2(out) # 2nd conv layer + out = out + identity # add residual + return out + +class DeconvLayer(nn.Module): + def __init__(self, in_channels, out_channels, kernel_size, stride, output_padding, norm="instance"): + super(DeconvLayer, self).__init__() + + # Transposed Convolution + padding_size = kernel_size // 2 + self.conv_transpose = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride, padding_size, output_padding) + + # Normalization Layers + self.norm_type = norm + if (norm=="instance"): + self.norm_layer = nn.InstanceNorm2d(out_channels, affine=True) + elif (norm=="batch"): + self.norm_layer = nn.BatchNorm2d(out_channels, affine=True) + + def forward(self, x): + x = self.conv_transpose(x) + if (self.norm_type=="None"): + out = x + else: + out = self.norm_layer(x) + return out + + + +def train(): + # Seeds + torch.manual_seed(SEED) + torch.cuda.manual_seed(SEED) + np.random.seed(SEED) + random.seed(SEED) + + # Device + device = ("cuda" if torch.cuda.is_available() else "cpu") + print(f"Train.py: Device is {device}") + + # Dataset and Dataloader + transform = transforms.Compose([ + transforms.Resize(TRAIN_IMAGE_SIZE), + transforms.CenterCrop(TRAIN_IMAGE_SIZE), + transforms.ToTensor(), + transforms.Lambda(lambda x: x.mul(255)) + ]) + train_dataset = datasets.ImageFolder(DATASET_PATH, transform=transform) + train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True) + print(f"Train.py: train_loader length: {len(train_loader)}") + + # Load networks + TransformerNetwork = TransformerNetworkClass().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 = style_tensor.add(imagenet_neg_mean) + B, C, H, W = style_tensor.shape + style_features = VGG(style_tensor.expand([BATCH_SIZE, C, H, W])) + style_gram = {} + for key, value in style_features.items(): + style_gram[key] = gram(value) + + # Optimizer settings + optimizer = optim.AdamW(TransformerNetwork.parameters(), lr=ADAM_LR, fused=True, amsgrad=True) + + # Loss trackers + content_loss_history = [] + style_loss_history = [] + total_loss_history = [] + 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 + optimizer.zero_grad() + + # Generate images and get features + content_batch = content_batch[:,[2,1,0]].to(device) + generated_batch = TransformerNetwork(content_batch) + content_features = VGG(content_batch.add(imagenet_neg_mean)) + generated_features = VGG(generated_batch.add(imagenet_neg_mean)) + + # Content Loss + MSELoss = nn.MSELoss().to(device) + content_loss = CONTENT_WEIGHT * MSELoss(generated_features['relu2_2'], content_features['relu2_2']) + batch_content_loss_sum += content_loss + + # Style Loss + style_loss = 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 + batch_style_loss_sum += style_loss.item() + + # Total Loss + total_loss = content_loss + style_loss + batch_total_loss_sum += total_loss.item() + + # Backprop and Weight Update + total_loss.backward() + optimizer.step() + + # Save Model and Print Losses + if (((batch_count-1)%SAVE_MODEL_EVERY == 0) or (batch_count==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("Time elapsed:\t{} seconds".format(time.time()-start_time)) + + # Save Model + 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" + 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) + + # 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 + print("Done Training the Transformer Network!") + print("Training Time: {} seconds".format(stop_time-start_time)) + print("========Content Loss========") + print(content_loss_history) + print("========Style Loss========") + print(style_loss_history) + 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()