diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/colorizator.py b/colorizator.py new file mode 100644 index 0000000..03a40e9 --- /dev/null +++ b/colorizator.py @@ -0,0 +1,68 @@ +import torch +from torchvision.transforms import ToTensor +import numpy as np + +from networks.colorizer import Colorizer +from denoising.denoiser import FFDNetDenoiser +from utils.utils import resize_pad, tile_process + + +class MangaColorizator: + def __init__(self, config): + if config.device == 'cuda' and not torch.cuda.is_available(): + print("[-] CUDA not available, using CPU.") + self.device = 'cpu' + else: + self.device = config.device + + self.tile_size = config.colorizer_tile_size + self.tile_pad = config.tile_pad + + self.model = Colorizer().to(self.device) + state_dict = torch.load(config.colorizer_path, map_location=self.device) + self.model.generator.load_state_dict(state_dict) + self.model = self.model.eval() + + self.denoiser = FFDNetDenoiser(self.device) + + self.current_image = None + self.current_hint = None + self.current_pad = None + + self.scale = 1 + + + def set_image(self, image, size=576, transform = ToTensor()): + if size % 32 != 0: + raise RuntimeError("[-] Size is not divisible by 32") + + image, self.current_pad = resize_pad(image, size) + self.current_image = transform(image).unsqueeze(0).to(self.device) + self.current_hint = torch.zeros(1, 4, self.current_image.shape[2], self.current_image.shape[3])\ + .float().to(self.device) + + def update_hint(self, hint, mask): + if issubclass(hint.dtype.type, np.integer): + hint = hint.astype('float32') / 255 + + hint = (hint - 0.5) / 0.5 + hint = torch.FloatTensor(hint).permute(2, 0, 1) + mask = torch.FloatTensor(np.expand_dims(mask, 0)) + + self.current_hint = torch.cat([hint * mask, mask], 0).unsqueeze(0).to(self.device) + + def colorize(self): + with torch.no_grad(): + img = torch.cat([self.current_image, self.current_hint], 1) + + if self.tile_size > 0: + fake_color = tile_process(self.model, img, self.scale, self.tile_size, self.tile_pad) + else: + fake_color, _ = self.model(img) + result = fake_color[0].detach().permute(1, 2, 0) * 0.5 + 0.5 + if self.current_pad[0] != 0: + result = result[:-self.current_pad[0]] + if self.current_pad[1] != 0: + result = result[:, :-self.current_pad[1]] + + return (result.detach().cpu().numpy() * 255.0).round().astype(np.uint8) \ No newline at end of file diff --git a/denoisator.py b/denoisator.py new file mode 100644 index 0000000..3f5464a --- /dev/null +++ b/denoisator.py @@ -0,0 +1,19 @@ +import torch + +from denoising.denoiser import FFDNetDenoiser + + +class MangaDenoiser: + def __init__(self, config): + if config.device == 'cuda' and not torch.cuda.is_available(): + print("[-] CUDA not available, using CPU.") + self.device = 'cpu' + else: + self.device = config.device + + self.model = FFDNetDenoiser(self.device) + + + def denoise(self, image, sigma=25): + with torch.no_grad(): + return self.model.get_denoised_image(image, sigma=sigma) diff --git a/denoising/__init__.py b/denoising/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/denoising/__pycache__/__init__.cpython-311.pyc b/denoising/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000..19cb926 Binary files /dev/null and b/denoising/__pycache__/__init__.cpython-311.pyc differ diff --git a/denoising/__pycache__/denoiser.cpython-311.pyc b/denoising/__pycache__/denoiser.cpython-311.pyc new file mode 100644 index 0000000..31a8042 Binary files /dev/null and b/denoising/__pycache__/denoiser.cpython-311.pyc differ diff --git a/denoising/__pycache__/functions.cpython-311.pyc b/denoising/__pycache__/functions.cpython-311.pyc new file mode 100644 index 0000000..54c3b68 Binary files /dev/null and b/denoising/__pycache__/functions.cpython-311.pyc differ diff --git a/denoising/__pycache__/models.cpython-311.pyc b/denoising/__pycache__/models.cpython-311.pyc new file mode 100644 index 0000000..cbe7fef Binary files /dev/null and b/denoising/__pycache__/models.cpython-311.pyc differ diff --git a/denoising/__pycache__/utils.cpython-311.pyc b/denoising/__pycache__/utils.cpython-311.pyc new file mode 100644 index 0000000..2acae6a Binary files /dev/null and b/denoising/__pycache__/utils.cpython-311.pyc differ diff --git a/denoising/denoiser.py b/denoising/denoiser.py new file mode 100644 index 0000000..c2bc0e4 --- /dev/null +++ b/denoising/denoiser.py @@ -0,0 +1,117 @@ +""" +Denoise an image with the FFDNet denoising method + +Copyright (C) 2018, Matias Tassano + +This program is free software: you can use, modify and/or +redistribute it under the terms of the GNU General Public +License as published by the Free Software Foundation, either +version 3 of the License, or (at your option) any later +version. You should have received a copy of this license along +this program. If not, see . +""" +import os +import argparse +import time + + +import numpy as np +import cv2 +import torch +import torch.nn as nn +from torch.autograd import Variable +from .models import FFDNet +from .utils import normalize, variable_to_cv2_image, remove_dataparallel_wrapper, is_rgb + +class FFDNetDenoiser: + def __init__(self, _device, _sigma = 25, _weights_dir = 'denoising/models/', _in_ch = 3): + self.sigma = _sigma / 255 + self.weights_dir = _weights_dir + self.channels = _in_ch + self.device = _device + + self.model = FFDNet(num_input_channels = _in_ch) + self.load_weights() + self.model.eval() + + + def load_weights(self): + weights_name = 'net_rgb.pth' if self.channels == 3 else 'net_gray.pth' + weights_path = os.path.join(self.weights_dir, weights_name) + if self.device == 'cuda': + state_dict = torch.load(weights_path, map_location=torch.device('cpu')) + device_ids = [0] + self.model = nn.DataParallel(self.model, device_ids=device_ids).cuda() + else: + state_dict = torch.load(weights_path, map_location='cpu') + # CPU mode: remove the DataParallel wrapper + state_dict = remove_dataparallel_wrapper(state_dict) + self.model.load_state_dict(state_dict) + + def get_denoised_image(self, imorig, sigma = None): + + if sigma is not None: + cur_sigma = sigma / 255 + else: + cur_sigma = self.sigma + + if len(imorig.shape) < 3 or imorig.shape[2] == 1: + imorig = np.repeat(np.expand_dims(imorig, 2), 3, 2) + + imorig = imorig[..., :3] + + if (max(imorig.shape[0], imorig.shape[1]) > 1200): + ratio = max(imorig.shape[0], imorig.shape[1]) / 1200 + imorig = cv2.resize(imorig, (int(imorig.shape[1] / ratio), int(imorig.shape[0] / ratio)), interpolation = cv2.INTER_AREA) + + imorig = imorig.transpose(2, 0, 1) + + if (imorig.max() > 1.2): + imorig = normalize(imorig) + imorig = np.expand_dims(imorig, 0) + + # Handle odd sizes + expanded_h = False + expanded_w = False + sh_im = imorig.shape + if sh_im[2]%2 == 1: + expanded_h = True + imorig = np.concatenate((imorig, imorig[:, :, -1, :][:, :, np.newaxis, :]), axis=2) + + if sh_im[3]%2 == 1: + expanded_w = True + imorig = np.concatenate((imorig, imorig[:, :, :, -1][:, :, :, np.newaxis]), axis=3) + + + imorig = torch.Tensor(imorig) + + + # Sets data type according to CPU or GPU modes + if self.device == 'cuda': + dtype = torch.cuda.FloatTensor + else: + dtype = torch.FloatTensor + + imnoisy = imorig.clone() + + + with torch.no_grad(): + imorig, imnoisy = imorig.type(dtype), imnoisy.type(dtype) + nsigma = torch.FloatTensor([cur_sigma]).type(dtype) + + + # Estimate noise and subtract it to the input image + im_noise_estim = self.model(imnoisy, nsigma) + outim = torch.clamp(imnoisy-im_noise_estim, 0., 1.) + + if expanded_h: + imorig = imorig[:, :, :-1, :] + outim = outim[:, :, :-1, :] + imnoisy = imnoisy[:, :, :-1, :] + + if expanded_w: + imorig = imorig[:, :, :, :-1] + outim = outim[:, :, :, :-1] + imnoisy = imnoisy[:, :, :, :-1] + + return variable_to_cv2_image(outim) diff --git a/denoising/functions.py b/denoising/functions.py new file mode 100644 index 0000000..954bc23 --- /dev/null +++ b/denoising/functions.py @@ -0,0 +1,101 @@ +""" +Functions implementing custom NN layers + +Copyright (C) 2018, Matias Tassano + +This program is free software: you can use, modify and/or +redistribute it under the terms of the GNU General Public +License as published by the Free Software Foundation, either +version 3 of the License, or (at your option) any later +version. You should have received a copy of this license along +this program. If not, see . +""" +import torch +from torch.autograd import Function, Variable + +def concatenate_input_noise_map(input, noise_sigma): + r"""Implements the first layer of FFDNet. This function returns a + torch.autograd.Variable composed of the concatenation of the downsampled + input image and the noise map. Each image of the batch of size CxHxW gets + converted to an array of size 4*CxH/2xW/2. Each of the pixels of the + non-overlapped 2x2 patches of the input image are placed in the new array + along the first dimension. + + Args: + input: batch containing CxHxW images + noise_sigma: the value of the pixels of the CxH/2xW/2 noise map + """ + # noise_sigma is a list of length batch_size + N, C, H, W = input.size() + dtype = input.type() + sca = 2 + sca2 = sca*sca + Cout = sca2*C + Hout = H//sca + Wout = W//sca + idxL = [[0, 0], [0, 1], [1, 0], [1, 1]] + + # Fill the downsampled image with zeros + if 'cuda' in dtype: + downsampledfeatures = torch.cuda.FloatTensor(N, Cout, Hout, Wout).fill_(0) + else: + downsampledfeatures = torch.FloatTensor(N, Cout, Hout, Wout).fill_(0) + + # Build the CxH/2xW/2 noise map + noise_map = noise_sigma.view(N, 1, 1, 1).repeat(1, C, Hout, Wout) + + # Populate output + for idx in range(sca2): + downsampledfeatures[:, idx:Cout:sca2, :, :] = \ + input[:, :, idxL[idx][0]::sca, idxL[idx][1]::sca] + + # concatenate de-interleaved mosaic with noise map + return torch.cat((noise_map, downsampledfeatures), 1) + +class UpSampleFeaturesFunction(Function): + r"""Extends PyTorch's modules by implementing a torch.autograd.Function. + This class implements the forward and backward methods of the last layer + of FFDNet. It basically performs the inverse of + concatenate_input_noise_map(): it converts each of the images of a + batch of size CxH/2xW/2 to images of size C/4xHxW + """ + @staticmethod + def forward(ctx, input): + N, Cin, Hin, Win = input.size() + dtype = input.type() + sca = 2 + sca2 = sca*sca + Cout = Cin//sca2 + Hout = Hin*sca + Wout = Win*sca + idxL = [[0, 0], [0, 1], [1, 0], [1, 1]] + + assert (Cin%sca2 == 0), 'Invalid input dimensions: number of channels should be divisible by 4' + + result = torch.zeros((N, Cout, Hout, Wout)).type(dtype) + for idx in range(sca2): + result[:, :, idxL[idx][0]::sca, idxL[idx][1]::sca] = input[:, idx:Cin:sca2, :, :] + + return result + + @staticmethod + def backward(ctx, grad_output): + N, Cg_out, Hg_out, Wg_out = grad_output.size() + dtype = grad_output.data.type() + sca = 2 + sca2 = sca*sca + Cg_in = sca2*Cg_out + Hg_in = Hg_out//sca + Wg_in = Wg_out//sca + idxL = [[0, 0], [0, 1], [1, 0], [1, 1]] + + # Build output + grad_input = torch.zeros((N, Cg_in, Hg_in, Wg_in)).type(dtype) + # Populate output + for idx in range(sca2): + grad_input[:, idx:Cg_in:sca2, :, :] = grad_output.data[:, :, idxL[idx][0]::sca, idxL[idx][1]::sca] + + return Variable(grad_input) + +# Alias functions +upsamplefeatures = UpSampleFeaturesFunction.apply diff --git a/denoising/models.py b/denoising/models.py new file mode 100644 index 0000000..ab50648 --- /dev/null +++ b/denoising/models.py @@ -0,0 +1,100 @@ +""" +Definition of the FFDNet model and its custom layers + +Copyright (C) 2018, Matias Tassano + +This program is free software: you can use, modify and/or +redistribute it under the terms of the GNU General Public +License as published by the Free Software Foundation, either +version 3 of the License, or (at your option) any later +version. You should have received a copy of this license along +this program. If not, see . +""" +import torch.nn as nn +from torch.autograd import Variable +import denoising.functions as functions + +class UpSampleFeatures(nn.Module): + r"""Implements the last layer of FFDNet + """ + def __init__(self): + super(UpSampleFeatures, self).__init__() + def forward(self, x): + return functions.upsamplefeatures(x) + +class IntermediateDnCNN(nn.Module): + r"""Implements the middel part of the FFDNet architecture, which + is basically a DnCNN net + """ + def __init__(self, input_features, middle_features, num_conv_layers): + super(IntermediateDnCNN, self).__init__() + self.kernel_size = 3 + self.padding = 1 + self.input_features = input_features + self.num_conv_layers = num_conv_layers + self.middle_features = middle_features + if self.input_features == 5: + self.output_features = 4 #Grayscale image + elif self.input_features == 15: + self.output_features = 12 #RGB image + else: + raise Exception('Invalid number of input features') + + layers = [] + layers.append(nn.Conv2d(in_channels=self.input_features,\ + out_channels=self.middle_features,\ + kernel_size=self.kernel_size,\ + padding=self.padding,\ + bias=False)) + layers.append(nn.ReLU(inplace=True)) + for _ in range(self.num_conv_layers-2): + layers.append(nn.Conv2d(in_channels=self.middle_features,\ + out_channels=self.middle_features,\ + kernel_size=self.kernel_size,\ + padding=self.padding,\ + bias=False)) + layers.append(nn.BatchNorm2d(self.middle_features)) + layers.append(nn.ReLU(inplace=True)) + layers.append(nn.Conv2d(in_channels=self.middle_features,\ + out_channels=self.output_features,\ + kernel_size=self.kernel_size,\ + padding=self.padding,\ + bias=False)) + self.itermediate_dncnn = nn.Sequential(*layers) + def forward(self, x): + out = self.itermediate_dncnn(x) + return out + +class FFDNet(nn.Module): + r"""Implements the FFDNet architecture + """ + def __init__(self, num_input_channels): + super(FFDNet, self).__init__() + self.num_input_channels = num_input_channels + if self.num_input_channels == 1: + # Grayscale image + self.num_feature_maps = 64 + self.num_conv_layers = 15 + self.downsampled_channels = 5 + self.output_features = 4 + elif self.num_input_channels == 3: + # RGB image + self.num_feature_maps = 96 + self.num_conv_layers = 12 + self.downsampled_channels = 15 + self.output_features = 12 + else: + raise Exception('Invalid number of input features') + + self.intermediate_dncnn = IntermediateDnCNN(\ + input_features=self.downsampled_channels,\ + middle_features=self.num_feature_maps,\ + num_conv_layers=self.num_conv_layers) + self.upsamplefeatures = UpSampleFeatures() + + def forward(self, x, noise_sigma): + concat_noise_x = functions.concatenate_input_noise_map(x.data, noise_sigma.data) + concat_noise_x = Variable(concat_noise_x) + h_dncnn = self.intermediate_dncnn(concat_noise_x) + pred_noise = self.upsamplefeatures(h_dncnn) + return pred_noise diff --git a/denoising/models/net_rgb.pth b/denoising/models/net_rgb.pth new file mode 100644 index 0000000..1132f2f Binary files /dev/null and b/denoising/models/net_rgb.pth differ diff --git a/denoising/utils.py b/denoising/utils.py new file mode 100644 index 0000000..faa5016 --- /dev/null +++ b/denoising/utils.py @@ -0,0 +1,66 @@ +""" +Different utilities such as orthogonalization of weights, initialization of +loggers, etc + +Copyright (C) 2018, Matias Tassano + +This program is free software: you can use, modify and/or +redistribute it under the terms of the GNU General Public +License as published by the Free Software Foundation, either +version 3 of the License, or (at your option) any later +version. You should have received a copy of this license along +this program. If not, see . +""" +import numpy as np +import cv2 + + +def variable_to_cv2_image(varim): + r"""Converts a torch.autograd.Variable to an OpenCV image + + Args: + varim: a torch.autograd.Variable + """ + nchannels = varim.size()[1] + if nchannels == 1: + res = (varim.data.cpu().numpy()[0, 0, :]*255.).clip(0, 255).astype(np.uint8) + elif nchannels == 3: + res = varim.data.cpu().numpy()[0] + res = cv2.cvtColor(res.transpose(1, 2, 0), cv2.COLOR_RGB2BGR) + res = (res*255.).clip(0, 255).astype(np.uint8) + else: + raise Exception('Number of color channels not supported') + return res + + +def normalize(data): + return np.float32(data/255.) + +def remove_dataparallel_wrapper(state_dict): + r"""Converts a DataParallel model to a normal one by removing the "module." + wrapper in the module dictionary + + Args: + state_dict: a torch.nn.DataParallel state dictionary + """ + from collections import OrderedDict + + new_state_dict = OrderedDict() + for k, vl in state_dict.items(): + name = k[7:] # remove 'module.' of DataParallel + new_state_dict[name] = vl + + return new_state_dict + +def is_rgb(im_path): + r""" Returns True if the image in im_path is an RGB image + """ + from skimage.io import imread + rgb = False + im = imread(im_path) + if (len(im.shape) == 3): + if not(np.allclose(im[...,0], im[...,1]) and np.allclose(im[...,2], im[...,1])): + rgb = True + print("rgb: {}".format(rgb)) + print("im shape: {}".format(im.shape)) + return rgb diff --git a/inference.py b/inference.py new file mode 100644 index 0000000..660486e --- /dev/null +++ b/inference.py @@ -0,0 +1,80 @@ +import argparse +import os +import time +import numpy as np +import PIL.Image as Image +from denoisator import MangaDenoiser +from colorizator import MangaColorizator +from upscalator import MangaUpscaler +from utils.utils import distance_from_grayscale, save_image, clear_torch_cache + +def process_image(image_path, output_folder, colorizer, upscaler, denoiser, config): + image_name = os.path.basename(image_path) + image = Image.open(image_path).convert("RGB") + image = np.array(image) + + coloredness = distance_from_grayscale(image) + if coloredness > 1: + print(f"[+] {image_name} is already colored, skipping.") + return + + if config.denoise: + print(f"[*] Denoising {image_name}...") + image = denoiser.denoise(image, config.denoise_sigma) + + if config.colorize: + print(f"[*] Colorizing {image_name}...") + colorizer.set_image((image.astype('float32') / 255), config.colorized_image_size) + image = colorizer.colorize() + + if config.upscale: + print(f"[*] Upscaling {image_name} by {config.upscale_factor}x...") + image = upscaler.upscale((image.astype('float32') / 255), config.upscale_factor) + + output_path = os.path.join(output_folder, image_name) + save_image(image, output_path) + print(f"[+] Processed {image_name} -> Saved to {output_path}") + +def main(): + parser = argparse.ArgumentParser(description="Batch Colorize Images") + parser.add_argument("--input_path", type=str, default="input", help="Folder containing images") + parser.add_argument("--output_path", type=str, default="output", help="Folder to save processed images") + + parser.add_argument('--device', choices=['cpu', 'cuda'], default='cuda', help='Device to use') + + parser.add_argument('--colorizer_path', default='networks/generator.zip') + parser.add_argument('--extractor_path', default='networks/extractor.pth') + parser.add_argument('--upscaler_path', default='networks/RealESRGAN_x4plus_anime_6B.pt') + parser.add_argument('--upscaler_type', choices=['ESRGAN', 'GigaGAN'], default='ESRGAN') + + parser.add_argument('--no-upscale', dest='upscale', action='store_false', default=True, help='Disable upscaling') + parser.add_argument('--no-colorize', dest='colorize', action='store_false', default=True, + help='Disable colorization') + parser.add_argument('--no-denoise', dest='denoise', action='store_false', default=True, help='Disable denoiser') + parser.add_argument('--upscale_factor', choices=[2, 4], default=4, type=int, help='Upscale by x2 or x4') + parser.add_argument('--denoise_sigma', default=25, type=int, help='How much noise to expect from the image') + + config = parser.parse_args() + os.makedirs(config.output_path, exist_ok=True) + + config.upscaler_tile_size = 256 + config.colorizer_tile_size = 0 + config.tile_pad = 8 + config.colorized_image_size = 576 # Width + + colorizer = MangaColorizator(config) if config.colorize else None + upscaler = MangaUpscaler(config) if config.upscale else None + denoiser = MangaDenoiser(config) if config.denoise else None + print("[+] Components initialized") + + images = [f for f in os.listdir(config.input_path) if f.lower().endswith((".png", ".jpg", ".jpeg", ".webp"))] + for img in images: + process_image(os.path.join(config.input_path, img), config.output_path, colorizer, upscaler, denoiser, config) + print("[+] Batch processing complete") + + clear_torch_cache() + print("[+] Components released") + + +if __name__ == "__main__": + main() diff --git a/launch.bat b/launch.bat new file mode 100644 index 0000000..ebac7df --- /dev/null +++ b/launch.bat @@ -0,0 +1,46 @@ +@echo off +echo ======================================= +echo Manga Colorize Launcher +echo ======================================= + +cd /d "C:\Users\Nighthawk\Desktop\manga_colorize" + +echo. +echo [1/3] Activating virtual environment... +call .venv\Scripts\activate +if %errorlevel% neq 0 ( + echo ERROR: Failed to activate virtual environment. + pause + exit /b +) + +echo. +echo [2/3] Installing dependencies from requirements.txt... +uv pip install -r requirements.txt +if %errorlevel% neq 0 ( + echo ERROR: Failed to install requirements. + pause + exit /b +) + +echo. +echo [3/3] Installing PyTorch for CUDA... +uv pip install torch torchvision --index-url https://download.pytorch.org/whl/cu126 +if %errorlevel% neq 0 ( + echo ERROR: Failed to install PyTorch. + pause + exit /b +) + +echo. +echo ======================================= +echo Launching Application... +echo ======================================= +echo. + +REM --- THIS IS THE KEY FIX --- +python main.py + +echo. +echo Application has closed. Press any key to exit the window. +pause \ No newline at end of file diff --git a/main.py b/main.py new file mode 100644 index 0000000..4007eb1 --- /dev/null +++ b/main.py @@ -0,0 +1,784 @@ +# main.py +import tkinter as tk +from tkinter import filedialog, messagebox, ttk +from PIL import Image, ImageTk +import threading +import os +import sys +import numpy as np +import argparse +import warnings +import platform +import shutil +import subprocess +import queue +import time +import cv2 + +def _prune_sys_path_configs(): + # Remove entries from sys.path that contain a top-level config.py, + # so cv2 won't accidentally exec it. + pruned = [] + for p in list(sys.path): + try: + base = p or os.getcwd() + if os.path.isfile(os.path.join(base, "config.py")): + # Drop this path entry + continue + except Exception: + pass + pruned.append(p) + sys.path[:] = pruned + +_prune_sys_path_configs() + +# --- THEME IMPORTS --- +import sv_ttk +import darkdetect +if platform.system() == "Windows": + try: + import pywinstyles + except ImportError: + pywinstyles = None + +# --- SUPPRESSION --- +try: + from torch.serialization import SourceChangeWarning + warnings.filterwarnings("ignore", category=SourceChangeWarning) +except ImportError: + pass + +# --- APP IMPORTS (Refactored) --- +# <-- FIX: Removed 'backend.' from imports to match file structure +from colorizator import MangaColorizator +from denoisator import MangaDenoiser +from upscalator import MangaUpscaler +from utils.utils import save_image + + +# ========================= +# Theme Manager +# ========================= +class ThemeManager: + def __init__(self, root: tk.Tk): + self.root = root + + def apply_windows_titlebar_theme(self): + if platform.system() != "Windows" or pywinstyles is None: + return + theme = sv_ttk.get_theme() + version = sys.getwindowsversion() + try: + if version.major == 10 and version.build >= 22000: + header_color = "#1c1c1c" if theme == "dark" else "#fafafa" + pywinstyles.change_header_color(self.root, header_color) + elif version.major == 10: + pywinstyles.apply_style(self.root, "dark" if theme == "dark" else "normal") + # little alpha flip to refresh caption colors + self.root.wm_attributes("-alpha", 0.99) + self.root.wm_attributes("-alpha", 1.0) + except Exception: + pass + + def toggle_theme(self): + current = sv_ttk.get_theme() + sv_ttk.set_theme("light" if current == "dark" else "dark") + self.apply_windows_titlebar_theme() + + def set_initial(self, theme_name: str): + sv_ttk.set_theme(theme_name) + self.apply_windows_titlebar_theme() + + +# ========================= +# Processor Pipeline +# ========================= +class ProcessorPipeline: + def __init__(self, config): + self.config = config + self.colorizer = MangaColorizator(config) if config.colorize else None + self.upscaler = MangaUpscaler(config) if config.upscale else None + self.denoiser = MangaDenoiser(config) if config.denoise else None + + def process(self, image_np: np.ndarray, + do_denoise: bool, do_colorize: bool, do_upscale: bool) -> np.ndarray: + out = image_np + if do_denoise and self.denoiser: + # Denoiser returns a BGR image, which needs conversion for subsequent steps + out = self.denoiser.denoise(out, self.config.denoise_sigma) + # Ensure RGB for colorizer (denoiser output is BGR) + out = cv2.cvtColor(out, cv2.COLOR_BGR2RGB) + + if do_colorize and self.colorizer: + self.colorizer.set_image((out.astype('float32') / 255.0), self.config.colorized_image_size) + out = self.colorizer.colorize() + + if do_upscale and self.upscaler: + out = self.upscaler.upscale((out.astype('float32') / 255.0), self.config.upscale_factor) + + return out + + +# ========================= +# Base Tab +# ========================= +class BaseTab(ttk.Frame): + def __init__(self, parent, app): + super().__init__(parent) + self.app = app # provides access to pipeline, theme, etc. + + +# ========================= +# Single Image Tab +# ========================= +class SingleImageTab(BaseTab): + def __init__(self, parent, app): + super().__init__(parent, app) + self.original_image = None + self.processed_image = None + self.image_path = None + + self.denoise_var = tk.BooleanVar(value=True) + self.colorize_var = tk.BooleanVar(value=True) + self.upscale_var = tk.BooleanVar(value=True) + + self._build() + + def _build(self): + main_frame = ttk.Frame(self, padding=10) + main_frame.pack(fill=tk.BOTH, expand=True) + + # Controls + control = ttk.Frame(main_frame) + control.pack(side=tk.TOP, fill=tk.X, pady=5) + + self.btn_open = ttk.Button(control, text="Open Image", command=self.open_image) + self.btn_open.pack(side=tk.LEFT, padx=5) + + self.btn_process = ttk.Button(control, text="Process Image", + command=self.start_processing, state=tk.DISABLED) + self.btn_process.pack(side=tk.LEFT, padx=5) + + self.btn_save = ttk.Button(control, text="Save Image", + command=self.save_image, state=tk.DISABLED) + self.btn_save.pack(side=tk.LEFT, padx=5) + + # Options + options = ttk.LabelFrame(main_frame, text="Processing Options", padding=10) + options.pack(side=tk.TOP, fill=tk.X, pady=10) + ttk.Checkbutton(options, text="Denoise", variable=self.denoise_var).pack(side=tk.LEFT, padx=10) + ttk.Checkbutton(options, text="Colorize", variable=self.colorize_var).pack(side=tk.LEFT, padx=10) + ttk.Checkbutton(options, text="Upscale", variable=self.upscale_var).pack(side=tk.LEFT, padx=10) + + # Image Panels + image_frame = ttk.Frame(main_frame) + image_frame.pack(fill=tk.BOTH, expand=True, pady=10) + self.panel_original = ttk.Label(image_frame, text="Original Image", relief="groove", anchor="center") + self.panel_original.pack(side=tk.LEFT, fill=tk.BOTH, expand=True, padx=5) + self.panel_processed = ttk.Label(image_frame, text="Processed Image", relief="groove", anchor="center") + self.panel_processed.pack(side=tk.RIGHT, fill=tk.BOTH, expand=True, padx=5) + + def open_image(self): + path = filedialog.askopenfilename(filetypes=[("Image Files", "*.png *.jpg *.jpeg *.webp")]) + if not path: + return + self.image_path = path + self.original_image = Image.open(path).convert("RGB") + self._display_image(self.original_image, self.panel_original) + self.panel_processed.config(image='', text="Processed Image") + self.processed_image = None + self.app.set_status(f"Loaded: {os.path.basename(path)}") + self._update_buttons() + + def _display_image(self, img: Image.Image, panel: ttk.Label): + panel_w, panel_h = panel.winfo_width(), panel.winfo_height() + if panel_w < 2 or panel_h < 2: + panel_w, panel_h = 550, 550 + p = img.copy() + p.thumbnail((panel_w, panel_h)) + photo = ImageTk.PhotoImage(p) + panel.config(image=photo) + panel.image = photo + + def start_processing(self): + self._set_ui_state(tk.DISABLED) + self.app.set_status("Processing...") + threading.Thread(target=self._process_worker, daemon=True).start() + + def _process_worker(self): + try: + src_np = np.array(self.original_image) + out_np = self.app.pipeline.process( + src_np, self.denoise_var.get(), self.colorize_var.get(), self.upscale_var.get() + ) + self.processed_image = Image.fromarray(out_np) + self.after(0, self._on_done) + except Exception as e: + self.after(0, lambda: messagebox.showerror("Processing Error", f"An error occurred: {e}")) + self.after(0, lambda: self._set_ui_state(tk.NORMAL)) + self.after(0, lambda: self.app.set_status("Error during processing.")) + + + def _on_done(self): + self._display_image(self.processed_image, self.panel_processed) + self._set_ui_state(tk.NORMAL) + self.app.set_status("Processing complete.") + self._update_buttons() + + def save_image(self): + if not self.processed_image: + return + save_path = filedialog.asksaveasfilename( + defaultextension=".png", + filetypes=[("PNG", "*.png"), ("JPEG", "*.jpg"), ("WebP", "*.webp")] + ) + if not save_path: + return + try: + fmt = os.path.splitext(save_path)[1][1:].upper() + if fmt == "JPG": + fmt = "JPEG" + save_image(np.array(self.processed_image), save_path, format=fmt) + messagebox.showinfo("Success", f"Image saved to {save_path}") + except Exception as e: + messagebox.showerror("Error", f"Failed to save image: {e}") + + def _update_buttons(self): + self.btn_process.config(state=tk.NORMAL if self.image_path else tk.DISABLED) + self.btn_save.config(state=tk.NORMAL if self.processed_image is not None else tk.DISABLED) + + def _set_ui_state(self, state): + self.btn_open.config(state=state) + # Re-evaluate process/save button states based on image presence + self.btn_process.config(state=state if self.image_path else tk.DISABLED) + self.btn_save.config(state=state if self.processed_image is not None else tk.DISABLED) + + +# ========================= +# Batch Tab +# ========================= +class BatchTab(BaseTab): + def __init__(self, parent, app): + super().__init__(parent, app) + self.input_folder = None + self.output_folder = None + self.gallery_widgets = {} + self.thumbnail_refs = [] + + self.denoise_var = tk.BooleanVar(value=True) + self.colorize_var = tk.BooleanVar(value=True) + self.upscale_var = tk.BooleanVar(value=True) + + self._build() + + def _build(self): + main = ttk.Frame(self, padding=10) + main.pack(fill=tk.BOTH, expand=True) + + control = ttk.Frame(main) + control.pack(side=tk.TOP, fill=tk.X, pady=5) + + self.btn_folders = ttk.Button(control, text="Select Folders", command=self.select_folders) + self.btn_folders.pack(side=tk.LEFT, padx=5) + + self.btn_start = ttk.Button(control, text="Start Batch Process", + command=self.start_batch, state=tk.DISABLED) + self.btn_start.pack(side=tk.LEFT, padx=5) + + self.folder_status = ttk.Label(control, text="No folders selected.") + self.folder_status.pack(side=tk.LEFT, padx=10, fill=tk.X, expand=True) + + # Options + options = ttk.LabelFrame(main, text="Processing Options", padding=10) + options.pack(side=tk.TOP, fill=tk.X, pady=10) + ttk.Checkbutton(options, text="Denoise", variable=self.denoise_var).pack(side=tk.LEFT, padx=10) + ttk.Checkbutton(options, text="Colorize", variable=self.colorize_var).pack(side=tk.LEFT, padx=10) + ttk.Checkbutton(options, text="Upscale", variable=self.upscale_var).pack(side=tk.LEFT, padx=10) + + # Gallery + gallery = ttk.LabelFrame(main, text="Image Preview", padding=10) + gallery.pack(fill=tk.BOTH, expand=True, pady=10) + + self.canvas = tk.Canvas(gallery, highlightthickness=0) + self.scroll = ttk.Scrollbar(gallery, orient="vertical", command=self.canvas.yview) + self.scrollable = ttk.Frame(self.canvas) + + self.scrollable.bind("", + lambda e: self.canvas.configure(scrollregion=self.canvas.bbox("all"))) + self.canvas.create_window((0, 0), window=self.scrollable, anchor="nw") + self.canvas.configure(yscrollcommand=self.scroll.set) + self.canvas.pack(side="left", fill="both", expand=True) + self.scroll.pack(side="right", fill="y") + + self.progress = ttk.Progressbar(main, orient="horizontal", mode="determinate") + self.progress.pack(side=tk.BOTTOM, fill=tk.X, pady=5) + + def select_folders(self): + in_dir = filedialog.askdirectory(title="Select Input Folder") + if not in_dir: + return + out_dir = filedialog.askdirectory(title="Select Output Folder") + if not out_dir: + return + self.input_folder = in_dir + self.output_folder = out_dir + self.folder_status.config(text=f"In: ...{in_dir[-30:]} | Out: ...{out_dir[-30:]}") + self.btn_start.config(state=tk.NORMAL) + threading.Thread(target=self._load_thumbs, daemon=True).start() + + def _load_thumbs(self): + self._set_ui_state(tk.DISABLED) + self.app.set_status("Loading thumbnails...") + for w in self.scrollable.winfo_children(): + w.destroy() + self.thumbnail_refs.clear() + self.gallery_widgets.clear() + + if not self.input_folder: + self._set_ui_state(tk.NORMAL) + return + + files = [f for f in os.listdir(self.input_folder) + if f.lower().endswith((".png", ".jpg", ".jpeg", ".webp"))] + + for i, fname in enumerate(files): + try: + img = Image.open(os.path.join(self.input_folder, fname)) + img.thumbnail((128, 128)) + photo = ImageTk.PhotoImage(img) + self.thumbnail_refs.append(photo) + self.after(0, self._add_thumb, photo, fname, i) + except Exception as e: + print(f"Thumb fail {fname}: {e}") + + self.after(0, lambda: self.app.set_status(f"Loaded {len(files)} images. Ready for batch.")) + self.after(0, lambda: self._set_ui_state(tk.NORMAL)) + + def _add_thumb(self, photo, fname, idx): + lbl = ttk.Label(self.scrollable, image=photo, text=fname, compound="top") + r, c = divmod(idx, 5) + lbl.grid(row=r, column=c, padx=5, pady=5) + self.gallery_widgets[fname] = lbl + + def _update_thumb(self, fname, npimg): + try: + img = Image.fromarray(npimg) + img.thumbnail((128, 128)) + ph = ImageTk.PhotoImage(img) + w = self.gallery_widgets.get(fname) + if w: + w.config(image=ph) + w.image = ph + except Exception as e: + print(f"Update thumb fail {fname}: {e}") + + def start_batch(self): + self._set_ui_state(tk.DISABLED) + threading.Thread(target=self._batch_worker, daemon=True).start() + + def _batch_worker(self): + try: + files = [f for f in os.listdir(self.input_folder) + if f.lower().endswith((".png", ".jpg", ".jpeg", ".webp"))] + total = len(files) + self.after(0, lambda: self.progress.config(maximum=total, value=0)) + for i, fname in enumerate(files, 1): + self.after(0, lambda i=i, f=fname, t=total: + self.app.set_status(f"Processing [{i}/{t}]: {f}")) + src = os.path.join(self.input_folder, fname) + npimg = np.array(Image.open(src).convert("RGB")) + out = self.app.pipeline.process( + npimg, self.denoise_var.get(), self.colorize_var.get(), self.upscale_var.get() + ) + # Save as PNG + base, _ = os.path.splitext(fname) + save_path = os.path.join(self.output_folder, f"{base}.png") + save_image(out, save_path, format="PNG") + self.after(0, self._update_thumb, fname, out) + self.after(0, self.progress.step) + self.after(0, lambda: messagebox.showinfo("Success", "Batch processing finished successfully.")) + self.after(0, lambda: self.app.set_status(f"Batch complete. Processed {total} files.")) + except Exception as e: + self.after(0, lambda: messagebox.showerror("Batch Processing Error", f"An error occurred: {e}")) + self.after(0, lambda: self.app.set_status("Error during batch processing.")) + finally: + self.after(0, lambda: self._set_ui_state(tk.NORMAL)) + self.after(0, lambda: self.progress.config(value=0)) + + def _set_ui_state(self, state): + self.btn_folders.config(state=state) + self.btn_start.config(state=state if self.input_folder else tk.DISABLED) + + +# # ========================= +# # Downloaders Tab (MangaDex, NHentai, possibly more.) +# # ========================= +# class DownloadersTab(BaseTab): + # def __init__(self, parent, app): + # super().__init__(parent, app) + # self.out_dir = tk.StringVar(value=os.path.abspath("downloads")) + # self._build() + + # self.proc_thread = None + # self.log_q = queue.Queue() + # self.stop_flag = threading.Event() + + # # ---------- UI ---------- + # def _build(self): + # container = ttk.Frame(self, padding=10) + # container.pack(fill=tk.BOTH, expand=True) + + # # Output row + # out_row = ttk.Frame(container) + # out_row.pack(fill=tk.X, pady=(0, 10)) + # ttk.Label(out_row, text="Output Folder:").pack(side=tk.LEFT) + # self.out_entry = ttk.Entry(out_row, textvariable=self.out_dir) + # self.out_entry.pack(side=tk.LEFT, fill=tk.X, expand=True, padx=8) + # ttk.Button(out_row, text="Browse", command=self._pick_out).pack(side=tk.LEFT) + + # # Notebook for providers + # self.provider_nb = ttk.Notebook(container) + # self.provider_nb.pack(fill=tk.BOTH, expand=True) + + # self._build_mangadex_tab() + # self._build_nhentai_tab() + + # # Log + actions + # bottom = ttk.Frame(container) + # bottom.pack(fill=tk.X, pady=(10, 0)) + + # ttk.Button(bottom, text="Open Folder", command=self._open_folder).pack(side=tk.LEFT) + # ttk.Button(bottom, text="Use as Batch Input", command=self._use_as_batch_input).pack(side=tk.LEFT, padx=8) + + # self.log = tk.Text(container, height=12, wrap="word") + # self.log.pack(fill=tk.BOTH, expand=False, pady=(10, 0)) + # self.log.configure(state="disabled") + + # def _build_mangadex_tab(self): + # tab = ttk.Frame(self.provider_nb, padding=10) + # self.provider_nb.add(tab, text="MangaDex") + + # ttk.Label(tab, text="MangaDex URL:").pack(anchor="w") + # self.md_url = tk.StringVar() + # ttk.Entry(tab, textvariable=self.md_url).pack(fill=tk.X, pady=5) + + # help_md = ( + # "Requirements:\n" + # "- pip install mangadex-downloader (optional extras: [optional])\n" + # "Usage examples:\n" + # " mangadex-dl \"\"\n" + # " python -m mangadex_downloader \"\"\n" + # ) + # ttk.Label(tab, text=help_md, justify="left").pack(anchor="w", pady=(0, 5)) + + # btns = ttk.Frame(tab) + # btns.pack(fill=tk.X, pady=5) + # self.btn_md_dl = ttk.Button(btns, text="Download", command=self._start_mangadex) + # self.btn_md_dl.pack(side=tk.LEFT) + + # def _build_nhentai_tab(self): + # tab = ttk.Frame(self.provider_nb, padding=10) + # self.provider_nb.add(tab, text="nhentai") + + # ttk.Label(tab, text="User-Agent (recommended):").pack(anchor="w") + # self.nh_ua = tk.StringVar() + # ttk.Entry(tab, textvariable=self.nh_ua).pack(fill=tk.X, pady=3) + + # ttk.Label(tab, text="Cookie (csrftoken=...; sessionid=...; cf_clearance=...):").pack(anchor="w") + # self.nh_cookie = tk.StringVar() + # ttk.Entry(tab, textvariable=self.nh_cookie).pack(fill=tk.X, pady=3) + + # ttk.Label(tab, text="Mode:").pack(anchor="w", pady=(6, 0)) + # self.nh_mode = tk.StringVar(value="id") + # mode_row = ttk.Frame(tab); mode_row.pack(anchor="w", pady=2) + # ttk.Radiobutton(mode_row, text="IDs", variable=self.nh_mode, value="id").pack(side=tk.LEFT) + # ttk.Radiobutton(mode_row, text="Search", variable=self.nh_mode, value="search").pack(side=tk.LEFT) + # ttk.Radiobutton(mode_row, text="Favorites (login cookie required)", variable=self.nh_mode, value="favorites").pack(side=tk.LEFT) + + # self.nh_ids = tk.StringVar() + # self.nh_query = tk.StringVar() + # self.nh_page = tk.StringVar(value="1") + # self.nh_download = tk.BooleanVar(value=True) + # self.nh_cbz = tk.BooleanVar(value=False) + # self.nh_pdf = tk.BooleanVar(value=False) + # self.nh_delay = tk.StringVar(value="0") + + # # IDs + # ids_row = ttk.Frame(tab); ids_row.pack(fill=tk.X, pady=3) + # ttk.Label(ids_row, text="IDs (space separated):").pack(side=tk.LEFT) + # ttk.Entry(ids_row, textvariable=self.nh_ids).pack(side=tk.LEFT, fill=tk.X, expand=True, padx=6) + + # # Search + # srch_row = ttk.Frame(tab); srch_row.pack(fill=tk.X, pady=3) + # ttk.Label(srch_row, text="Search:").pack(side=tk.LEFT) + # ttk.Entry(srch_row, textvariable=self.nh_query).pack(side=tk.LEFT, fill=tk.X, expand=True, padx=6) + # ttk.Label(srch_row, text="Page:").pack(side=tk.LEFT, padx=(8, 2)) + # ttk.Entry(srch_row, width=6, textvariable=self.nh_page).pack(side=tk.LEFT) + + # # Options + # opt_row = ttk.Frame(tab); opt_row.pack(fill=tk.X, pady=5) + # ttk.Checkbutton(opt_row, text="Download", variable=self.nh_download).pack(side=tk.LEFT) + # ttk.Checkbutton(opt_row, text="CBZ", variable=self.nh_cbz).pack(side=tk.LEFT, padx=6) + # ttk.Checkbutton(opt_row, text="PDF", variable=self.nh_pdf).pack(side=tk.LEFT, padx=6) + # ttk.Label(opt_row, text="Delay (s):").pack(side=tk.LEFT, padx=(10, 2)) + # ttk.Entry(opt_row, width=6, textvariable=self.nh_delay).pack(side=tk.LEFT) + + # help_nh = ( + # "Tips:\n" + # "- pip install nhentai\n" + # "- To bypass Cloudflare rate limits, set both --cookie and --useragent.\n" + # "- Use same IP/User-Agent as when the cookie was obtained.\n" + # ) + # ttk.Label(tab, text=help_nh, justify="left").pack(anchor="w", pady=(4, 6)) + + # btns = ttk.Frame(tab); btns.pack(fill=tk.X, pady=5) + # self.btn_nh_go = ttk.Button(btns, text="Run nhentai", command=self._start_nhentai) + # self.btn_nh_go.pack(side=tk.LEFT) + + # # ---------- Actions ---------- + # def _pick_out(self): + # d = filedialog.askdirectory(title="Select Output Folder") + # if d: + # self.out_dir.set(d) + + # def _open_folder(self): + # path = self.out_dir.get() + # if not os.path.isdir(path): + # messagebox.showerror("Error", "Output folder does not exist.") + # return + # if platform.system() == "Windows": + # os.startfile(path) + # elif platform.system() == "Darwin": + # subprocess.Popen(["open", path]) + # else: + # subprocess.Popen(["xdg-open", path]) + + # def _use_as_batch_input(self): + # # Tell App to switch to Batch tab and set input folder + # self.app.use_folder_as_batch_input(self.out_dir.get()) + + # # ---------- Logging ---------- + # def _log(self, text: str): + # self.log.configure(state="normal") + # self.log.insert("end", text) + # self.log.see("end") + # self.log.configure(state="disabled") + + # def _pump_logq(self): + # try: + # while True: + # line = self.log_q.get_nowait() + # self._log(line) + # except queue.Empty: + # pass + # if not self.stop_flag.is_set(): + # self.after(100, self._pump_logq) + + # # ---------- Runner helpers ---------- + # def _which_or_module(self, cli_names, module_cmd): + # """ + # Return a list representing the command to execute. Try CLIs in order, + # else return ['python', '-m', module_cmd]. + # """ + # for name in cli_names: + # if shutil.which(name): + # return [name] + # # Fallback to python -m module + # py = shutil.which("python3") or shutil.which("python") or sys.executable + # return [py, "-m", module_cmd] + + # def _start_mangadex(self): + # url = self.md_url.get().strip() + # if not url: + # messagebox.showerror("Error", "Please enter a MangaDex URL.") + # return + # out = self.out_dir.get() + # os.makedirs(out, exist_ok=True) + + # base = self._which_or_module(["mangadex-dl", "mangadex-downloader"], "mangadex_downloader") + # cmd = base + [url, "-o", out] + + # self._run_cmd_threaded("MangaDex", cmd) + + # def _start_nhentai(self): + # out = self.out_dir.get() + # os.makedirs(out, exist_ok=True) + + # base = self._which_or_module(["nhentai"], "nhentai") + # cmd = base + ["-o", out] + + # ua = self.nh_ua.get().strip() + # ck = self.nh_cookie.get().strip() + # if ua: + # cmd += ["--useragent", ua] + # if ck: + # cmd += ["--cookie", ck] + + # mode = self.nh_mode.get() + # if mode == "id": + # ids = self.nh_ids.get().strip() + # if not ids: + # messagebox.showerror("Error", "Enter at least one ID.") + # return + # cmd += ["--id"] + ids.split() + # elif mode == "search": + # q = self.nh_query.get().strip() + # if not q: + # messagebox.showerror("Error", "Enter a search query.") + # return + # page = self.nh_page.get().strip() or "1" + # cmd += ["--search", q, "--page", page] + # if self.nh_download.get(): + # cmd.append("--download") + # else: # favorites + # cmd += ["--favorites"] + # if self.nh_download.get(): + # cmd.append("--download") + # d = self.nh_delay.get().strip() + # if d and d.isdigit(): + # cmd += ["--delay", d] + + # if self.nh_cbz.get(): + # cmd.append("--cbz") + # if self.nh_pdf.get(): + # cmd.append("--pdf") + + # self._run_cmd_threaded("nhentai", cmd) + + # def _run_cmd_threaded(self, label, cmd): + # if self.proc_thread and self.proc_thread.is_alive(): + # messagebox.showwarning("Busy", "A download is already in progress.") + # return + + # self.stop_flag.clear() + # self._log(f"\n[{label}] Running: {' '.join(cmd)}\n") + # self.after(100, self._pump_logq) + + # def worker(): + # try: + # proc = subprocess.Popen( + # cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1, + # creationflags=subprocess.CREATE_NO_WINDOW if platform.system() == "Windows" else 0 + # ) + # for line in iter(proc.stdout.readline, ''): + # self.log_q.put(line) + # proc.stdout.close() + # rc = proc.wait() + # self.log_q.put(f"[{label}] Finished with code {rc}\n") + # if rc == 0: + # self.log_q.put("Done. You can Open Folder or Use as Batch Input.\n") + # except FileNotFoundError: + # self.log_q.put(f"[{label}] Command not found. Is the tool installed?\n") + # except Exception as e: + # self.log_q.put(f"[{label}] Error: {e}\n") + # finally: + # self.stop_flag.set() + + # self.proc_thread = threading.Thread(target=worker, daemon=True) + # self.proc_thread.start() + + +# ========================= +# App +# ========================= +class App: + def __init__(self, root: tk.Tk, config): + self.root = root + self.root.title("Manga Image Processor") + self.root.geometry("1200x800") + self.config = config + + # Theme + self.theme = ThemeManager(root) + + # Pipeline + try: + self.pipeline = ProcessorPipeline(config) + print("[+] Components initialized successfully.") + except Exception as e: + messagebox.showerror("Initialization Error", f"Failed to initialize components: {e}") + root.destroy() + return + + # Header + header = ttk.Frame(root, padding=(10, 10, 10, 0)) + header.pack(fill=tk.X) + ttk.Label(header, text="").pack(side=tk.LEFT, expand=True) + ttk.Button(header, text="Toggle Theme", command=self.theme.toggle_theme).pack(side=tk.RIGHT) + + # Tabs + self.nb = ttk.Notebook(root) + self.nb.pack(fill=tk.BOTH, expand=True, padx=10, pady=(0, 10)) + + self.tab_single = SingleImageTab(self.nb, self) + self.tab_batch = BatchTab(self.nb, self) + # self.tab_dl = DownloadersTab(self.nb, self) + + self.nb.add(self.tab_single, text="Single Image Processing") + self.nb.add(self.tab_batch, text="Batch Folder Processing") + # self.nb.add(self.tab_dl, text="Downloaders") + + # Status + self.status = ttk.Label(root, text="Ready", padding=(10, 5)) + self.status.pack(side=tk.BOTTOM, fill=tk.X) + + def set_status(self, text: str): + self.status.config(text=text) + + def use_folder_as_batch_input(self, folder: str): + if not os.path.isdir(folder): + messagebox.showerror("Error", "Output folder does not exist.") + return + # Switch to Batch tab & set input/output quickly + self.nb.select(self.tab_batch) + self.tab_batch.input_folder = folder + # If user hasn’t chosen output, default to "_processed" + default_out = folder.rstrip("/\\") + "_processed" + os.makedirs(default_out, exist_ok=True) + self.tab_batch.output_folder = default_out + self.tab_batch.folder_status.config( + text=f"In: ...{folder[-30:]} | Out: ...{default_out[-30:]}") + self.tab_batch.btn_start.config(state=tk.NORMAL) + threading.Thread(target=self.tab_batch._load_thumbs, daemon=True).start() + + +# ========================= +# Entrypoint +# ========================= +def parse_args(): + parser = argparse.ArgumentParser(description="GUI for Manga Image Processor") + parser.add_argument('--device', choices=['cpu', 'cuda'], default='cuda', help='Device to use') + + # --- Paths (Refactored) --- + # <-- FIX: Removed 'backend' from default paths + parser.add_argument('--colorizer_path', default='networks/generator.zip') + parser.add_argument('--extractor_path', default='networks/extractor.pth') + parser.add_argument('--upscaler_path', default='networks/RealESRGAN_x4plus_anime_6B.pt') + + parser.add_argument('--upscaler_type', choices=['ESRGAN', 'GigaGAN'], default='ESRGAN') + parser.add_argument('--no-upscale', dest='upscale', action='store_false', default=True) + parser.add_argument('--no-colorize', dest='colorize', action='store_false', default=True) + parser.add_argument('--no-denoise', dest='denoise', action='store_false', default=True) + parser.add_argument('--upscale_factor', choices=[2, 4], default=4, type=int) + parser.add_argument('--denoise_sigma', default=25, type=int) + + # Extra runtime tunables you already set later: + args = parser.parse_args() + args.upscaler_tile_size = 256 + args.colorizer_tile_size = 0 + args.tile_pad = 8 + args.colorized_image_size = 576 + return args + + +def main(): + config = parse_args() + root = tk.Tk() + app = App(root, config) + initial_theme = "dark" if darkdetect.isDark() else "light" + root.after(10, lambda: app.theme.set_initial(initial_theme)) + root.mainloop() + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/networks/RRDBNet.py b/networks/RRDBNet.py new file mode 100644 index 0000000..f387422 --- /dev/null +++ b/networks/RRDBNet.py @@ -0,0 +1,159 @@ +#https://github.com/XPixelGroup/BasicSR/blob/master/basicsr/archs/rrdbnet_arch.py +#https://github.com/XPixelGroup/BasicSR/blob/master/basicsr/archs/arch_util.py +import torch +from torch import nn as nn +from torch.nn import functional as F + +#from .arch_util import make_layer, pixel_unshuffle + +def pixel_unshuffle(x, scale): + """ Pixel unshuffle. + Args: + x (Tensor): Input feature with shape (b, c, hh, hw). + scale (int): Downsample ratio. + Returns: + Tensor: the pixel unshuffled feature. + """ + b, c, hh, hw = x.size() + out_channel = c * (scale**2) + assert hh % scale == 0 and hw % scale == 0 + h = hh // scale + w = hw // scale + x_view = x.view(b, c, h, scale, w, scale) + return x_view.permute(0, 1, 3, 5, 2, 4).reshape(b, out_channel, h, w) + +def make_layer(basic_block, num_basic_block, **kwarg): + """Make layers by stacking the same blocks. + Args: + basic_block (nn.module): nn.module class for basic block. + num_basic_block (int): number of blocks. + Returns: + nn.Sequential: Stacked blocks in nn.Sequential. + """ + layers = [] + for _ in range(num_basic_block): + layers.append(basic_block(**kwarg)) + return nn.Sequential(*layers) + +class ResidualDenseBlock(nn.Module): + """Residual Dense Block. + + Used in RRDB block in ESRGAN. + + Args: + num_feat (int): Channel number of intermediate features. + num_grow_ch (int): Channels for each growth. + """ + + def __init__(self, num_feat=64, num_grow_ch=32): + super(ResidualDenseBlock, self).__init__() + self.conv1 = nn.Conv2d(num_feat, num_grow_ch, 3, 1, 1) + self.conv2 = nn.Conv2d(num_feat + num_grow_ch, num_grow_ch, 3, 1, 1) + self.conv3 = nn.Conv2d(num_feat + 2 * num_grow_ch, num_grow_ch, 3, 1, 1) + self.conv4 = nn.Conv2d(num_feat + 3 * num_grow_ch, num_grow_ch, 3, 1, 1) + self.conv5 = nn.Conv2d(num_feat + 4 * num_grow_ch, num_feat, 3, 1, 1) + + self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True) + + # initialization + #default_init_weights([self.conv1, self.conv2, self.conv3, self.conv4, self.conv5], 0.1) + + def forward(self, x): + x1 = self.lrelu(self.conv1(x)) + x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1))) + x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1))) + x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1))) + x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1)) + # Emperically, we use 0.2 to scale the residual for better performance + return x5 * 0.2 + x + + +class RRDB(nn.Module): + """Residual in Residual Dense Block. + + Used in RRDB-Net in ESRGAN. + + Args: + num_feat (int): Channel number of intermediate features. + num_grow_ch (int): Channels for each growth. + """ + + def __init__(self, num_feat, num_grow_ch=32): + super(RRDB, self).__init__() + self.rdb1 = ResidualDenseBlock(num_feat, num_grow_ch) + self.rdb2 = ResidualDenseBlock(num_feat, num_grow_ch) + self.rdb3 = ResidualDenseBlock(num_feat, num_grow_ch) + + def forward(self, x): + out = self.rdb1(x) + out = self.rdb2(out) + out = self.rdb3(out) + # Emperically, we use 0.2 to scale the residual for better performance + return out * 0.2 + x + + + +class RRDBNet(nn.Module): + """Networks consisting of Residual in Residual Dense Block, which is used + in ESRGAN. + + ESRGAN: Enhanced Super-Resolution Generative Adversarial Networks. + + We extend ESRGAN for scale x2 and scale x1. + Note: This is one option for scale 1, scale 2 in RRDBNet. + We first employ the pixel-unshuffle (an inverse operation of pixelshuffle to reduce the spatial size + and enlarge the channel size before feeding inputs into the main ESRGAN architecture. + + Args: + num_in_ch (int): Channel number of inputs. + num_out_ch (int): Channel number of outputs. + num_feat (int): Channel number of intermediate features. + Default: 64 + num_block (int): Block number in the trunk network. Defaults: 23 + num_grow_ch (int): Channels for each growth. Default: 32. + """ + + def __init__(self, num_in_ch, num_out_ch, scale=4, num_feat=64, num_block=23, num_grow_ch=32): + super(RRDBNet, self).__init__() + self.scale = scale + if scale == 2: + num_in_ch = num_in_ch * 4 + elif scale == 1: + num_in_ch = num_in_ch * 16 + self.conv_first = nn.Conv2d(num_in_ch, num_feat, 3, 1, 1) + self.body = make_layer(RRDB, num_block, num_feat=num_feat, num_grow_ch=num_grow_ch) + self.conv_body = nn.Conv2d(num_feat, num_feat, 3, 1, 1) + # upsample + self.conv_up1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1) + self.conv_up2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1) + self.conv_hr = nn.Conv2d(num_feat, num_feat, 3, 1, 1) + self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1) + + self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True) + + def forward(self, x): + if self.scale == 2: + feat = pixel_unshuffle(x, scale=2) + elif self.scale == 1: + feat = pixel_unshuffle(x, scale=4) + else: + feat = x + feat = self.conv_first(feat) + body_feat = self.conv_body(self.body(feat)) + feat = feat + body_feat + # upsample + feat = self.lrelu(self.conv_up1(F.interpolate(feat, scale_factor=2, mode='nearest'))) + feat = self.lrelu(self.conv_up2(F.interpolate(feat, scale_factor=2, mode='nearest'))) + out = self.conv_last(self.lrelu(self.conv_hr(feat))) + return out + +class Upscaler(nn.Module): + def __init__(self, is_rgb=True): + super(Upscaler, self).__init__() + + channels = 3 if is_rgb else 1 + self.name = 'upscaler' + self.generator = RRDBNet(channels, channels) + + def forward(self, x): + return self.generator(x) \ No newline at end of file diff --git a/networks/RealESRGAN_x4plus_anime_6B.pt b/networks/RealESRGAN_x4plus_anime_6B.pt new file mode 100644 index 0000000..f2ae622 Binary files /dev/null and b/networks/RealESRGAN_x4plus_anime_6B.pt differ diff --git a/networks/__init__.py b/networks/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/networks/__pycache__/RRDBNet.cpython-311.pyc b/networks/__pycache__/RRDBNet.cpython-311.pyc new file mode 100644 index 0000000..4881edf Binary files /dev/null and b/networks/__pycache__/RRDBNet.cpython-311.pyc differ diff --git a/networks/__pycache__/__init__.cpython-311.pyc b/networks/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000..0c6913f Binary files /dev/null and b/networks/__pycache__/__init__.cpython-311.pyc differ diff --git a/networks/__pycache__/aura_sr.cpython-311.pyc b/networks/__pycache__/aura_sr.cpython-311.pyc new file mode 100644 index 0000000..6975b21 Binary files /dev/null and b/networks/__pycache__/aura_sr.cpython-311.pyc differ diff --git a/networks/__pycache__/colorizer.cpython-311.pyc b/networks/__pycache__/colorizer.cpython-311.pyc new file mode 100644 index 0000000..fa1fd15 Binary files /dev/null and b/networks/__pycache__/colorizer.cpython-311.pyc differ diff --git a/networks/__pycache__/extractor.cpython-311.pyc b/networks/__pycache__/extractor.cpython-311.pyc new file mode 100644 index 0000000..787814b Binary files /dev/null and b/networks/__pycache__/extractor.cpython-311.pyc differ diff --git a/networks/aura_sr.py b/networks/aura_sr.py new file mode 100644 index 0000000..6b6b413 --- /dev/null +++ b/networks/aura_sr.py @@ -0,0 +1,928 @@ +# AuraSR: GAN-based Super-Resolution for real-world, a reproduction of the GigaGAN* paper. Implementation is +# based on the unofficial lucidrains/gigagan-pytorch repository. Heavily modified from there. +# +# https://mingukkang.github.io/GigaGAN/ +from math import log2, ceil +from functools import partial +from typing import Any, Optional, List, Iterable + +import numpy as np +import torch +from torchvision import transforms +from PIL import Image +from torch import nn, einsum, Tensor +import torch.nn.functional as F + +from einops import rearrange, repeat, reduce +from einops.layers.torch import Rearrange +import math + + +def get_same_padding(size, kernel, dilation, stride): + return ((size - 1) * (stride - 1) + dilation * (kernel - 1)) // 2 + + +class AdaptiveConv2DMod(nn.Module): + def __init__( + self, + dim, + dim_out, + kernel, + *, + demod=True, + stride=1, + dilation=1, + eps=1e-8, + num_conv_kernels=1, # set this to be greater than 1 for adaptive + ): + super().__init__() + self.eps = eps + + self.dim_out = dim_out + + self.kernel = kernel + self.stride = stride + self.dilation = dilation + self.adaptive = num_conv_kernels > 1 + + self.weights = nn.Parameter( + torch.randn((num_conv_kernels, dim_out, dim, kernel, kernel)) + ) + + self.demod = demod + + nn.init.kaiming_normal_( + self.weights, a=0, mode="fan_in", nonlinearity="leaky_relu" + ) + + def forward( + self, fmap, mod: Optional[Tensor] = None, kernel_mod: Optional[Tensor] = None + ): + """ + notation + + b - batch + n - convs + o - output + i - input + k - kernel + """ + + b, h = fmap.shape[0], fmap.shape[-2] + + # account for feature map that has been expanded by the scale in the first dimension + # due to multiscale inputs and outputs + + if mod.shape[0] != b: + mod = repeat(mod, "b ... -> (s b) ...", s=b // mod.shape[0]) + + if exists(kernel_mod): + kernel_mod_has_el = kernel_mod.numel() > 0 + + assert self.adaptive or not kernel_mod_has_el + + if kernel_mod_has_el and kernel_mod.shape[0] != b: + kernel_mod = repeat( + kernel_mod, "b ... -> (s b) ...", s=b // kernel_mod.shape[0] + ) + + # prepare weights for modulation + + weights = self.weights + + if self.adaptive: + weights = repeat(weights, "... -> b ...", b=b) + + # determine an adaptive weight and 'select' the kernel to use with softmax + + assert exists(kernel_mod) and kernel_mod.numel() > 0 + + kernel_attn = kernel_mod.softmax(dim=-1) + kernel_attn = rearrange(kernel_attn, "b n -> b n 1 1 1 1") + + weights = reduce(weights * kernel_attn, "b n ... -> b ...", "sum") + + # do the modulation, demodulation, as done in stylegan2 + + mod = rearrange(mod, "b i -> b 1 i 1 1") + + weights = weights * (mod + 1) + + if self.demod: + inv_norm = ( + reduce(weights ** 2, "b o i k1 k2 -> b o 1 1 1", "sum") + .clamp(min=self.eps) + .rsqrt() + ) + weights = weights * inv_norm + + fmap = rearrange(fmap, "b c h w -> 1 (b c) h w") + + weights = rearrange(weights, "b o ... -> (b o) ...") + + padding = get_same_padding(h, self.kernel, self.dilation, self.stride) + fmap = F.conv2d(fmap, weights, padding=padding, groups=b) + + return rearrange(fmap, "1 (b o) ... -> b o ...", b=b) + + +class Attend(nn.Module): + def __init__(self, dropout=0.0, flash=False): + super().__init__() + self.dropout = dropout + self.attn_dropout = nn.Dropout(dropout) + self.scale = nn.Parameter(torch.randn(1)) + self.flash = flash + + def flash_attn(self, q, k, v): + q, k, v = map(lambda t: t.contiguous(), (q, k, v)) + out = F.scaled_dot_product_attention( + q, k, v, dropout_p=self.dropout if self.training else 0.0 + ) + return out + + def forward(self, q, k, v): + if self.flash: + return self.flash_attn(q, k, v) + + scale = q.shape[-1] ** -0.5 + + # similarity + sim = einsum("b h i d, b h j d -> b h i j", q, k) * scale + + # attention + attn = sim.softmax(dim=-1) + attn = self.attn_dropout(attn) + + # aggregate values + out = einsum("b h i j, b h j d -> b h i d", attn, v) + + return out + + +def exists(x): + return x is not None + + +def default(val, d): + if exists(val): + return val + return d() if callable(d) else d + + +def cast_tuple(t, length=1): + if isinstance(t, tuple): + return t + return (t,) * length + + +def identity(t, *args, **kwargs): + return t + + +def is_power_of_two(n): + return log2(n).is_integer() + + +def null_iterator(): + while True: + yield None + + +def Downsample(dim, dim_out=None): + return nn.Sequential( + Rearrange("b c (h p1) (w p2) -> b (c p1 p2) h w", p1=2, p2=2), + nn.Conv2d(dim * 4, default(dim_out, dim), 1), + ) + + +class RMSNorm(nn.Module): + def __init__(self, dim): + super().__init__() + self.g = nn.Parameter(torch.ones(1, dim, 1, 1)) + self.eps = 1e-4 + + def forward(self, x): + return F.normalize(x, dim=1) * self.g * (x.shape[1] ** 0.5) + + +# building block modules + + +class Block(nn.Module): + def __init__(self, dim, dim_out, groups=8, num_conv_kernels=0): + super().__init__() + self.proj = AdaptiveConv2DMod( + dim, dim_out, kernel=3, num_conv_kernels=num_conv_kernels + ) + self.kernel = 3 + self.dilation = 1 + self.stride = 1 + + self.act = nn.SiLU() + + def forward(self, x, conv_mods_iter: Optional[Iterable] = None): + conv_mods_iter = default(conv_mods_iter, null_iterator()) + + x = self.proj(x, mod=next(conv_mods_iter), kernel_mod=next(conv_mods_iter)) + + x = self.act(x) + return x + + +class ResnetBlock(nn.Module): + def __init__( + self, dim, dim_out, *, groups=8, num_conv_kernels=0, style_dims: List = [] + ): + super().__init__() + style_dims.extend([dim, num_conv_kernels, dim_out, num_conv_kernels]) + + self.block1 = Block( + dim, dim_out, groups=groups, num_conv_kernels=num_conv_kernels + ) + self.block2 = Block( + dim_out, dim_out, groups=groups, num_conv_kernels=num_conv_kernels + ) + self.res_conv = nn.Conv2d(dim, dim_out, 1) if dim != dim_out else nn.Identity() + + def forward(self, x, conv_mods_iter: Optional[Iterable] = None): + h = self.block1(x, conv_mods_iter=conv_mods_iter) + h = self.block2(h, conv_mods_iter=conv_mods_iter) + + return h + self.res_conv(x) + + +class LinearAttention(nn.Module): + def __init__(self, dim, heads=4, dim_head=32): + super().__init__() + self.scale = dim_head ** -0.5 + self.heads = heads + hidden_dim = dim_head * heads + + self.norm = RMSNorm(dim) + self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False) + + self.to_out = nn.Sequential(nn.Conv2d(hidden_dim, dim, 1), RMSNorm(dim)) + + def forward(self, x): + b, c, h, w = x.shape + + x = self.norm(x) + + qkv = self.to_qkv(x).chunk(3, dim=1) + q, k, v = map( + lambda t: rearrange(t, "b (h c) x y -> b h c (x y)", h=self.heads), qkv + ) + + q = q.softmax(dim=-2) + k = k.softmax(dim=-1) + + q = q * self.scale + + context = torch.einsum("b h d n, b h e n -> b h d e", k, v) + + out = torch.einsum("b h d e, b h d n -> b h e n", context, q) + out = rearrange(out, "b h c (x y) -> b (h c) x y", h=self.heads, x=h, y=w) + return self.to_out(out) + + +class Attention(nn.Module): + def __init__(self, dim, heads=4, dim_head=32, flash=False): + super().__init__() + self.heads = heads + hidden_dim = dim_head * heads + + self.norm = RMSNorm(dim) + + self.attend = Attend(flash=flash) + self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False) + self.to_out = nn.Conv2d(hidden_dim, dim, 1) + + def forward(self, x): + b, c, h, w = x.shape + x = self.norm(x) + qkv = self.to_qkv(x).chunk(3, dim=1) + + q, k, v = map( + lambda t: rearrange(t, "b (h c) x y -> b h (x y) c", h=self.heads), qkv + ) + + out = self.attend(q, k, v) + out = rearrange(out, "b h (x y) d -> b (h d) x y", x=h, y=w) + + return self.to_out(out) + + +# feedforward +def FeedForward(dim, mult=4): + return nn.Sequential( + RMSNorm(dim), + nn.Conv2d(dim, dim * mult, 1), + nn.GELU(), + nn.Conv2d(dim * mult, dim, 1), + ) + + +# transformers +class Transformer(nn.Module): + def __init__(self, dim, dim_head=64, heads=8, depth=1, flash_attn=True, ff_mult=4): + super().__init__() + self.layers = nn.ModuleList([]) + + for _ in range(depth): + self.layers.append( + nn.ModuleList( + [ + Attention( + dim=dim, dim_head=dim_head, heads=heads, flash=flash_attn + ), + FeedForward(dim=dim, mult=ff_mult), + ] + ) + ) + + def forward(self, x): + for attn, ff in self.layers: + x = attn(x) + x + x = ff(x) + x + + return x + + +class LinearTransformer(nn.Module): + def __init__(self, dim, dim_head=64, heads=8, depth=1, ff_mult=4): + super().__init__() + self.layers = nn.ModuleList([]) + + for _ in range(depth): + self.layers.append( + nn.ModuleList( + [ + LinearAttention(dim=dim, dim_head=dim_head, heads=heads), + FeedForward(dim=dim, mult=ff_mult), + ] + ) + ) + + def forward(self, x): + for attn, ff in self.layers: + x = attn(x) + x + x = ff(x) + x + + return x + + +class NearestNeighborhoodUpsample(nn.Module): + def __init__(self, dim, dim_out=None): + super().__init__() + dim_out = default(dim_out, dim) + self.conv = nn.Conv2d(dim, dim_out, kernel_size=3, stride=1, padding=1) + + def forward(self, x): + if x.shape[0] >= 64: + x = x.contiguous() + + x = F.interpolate(x, scale_factor=2.0, mode="nearest") + x = self.conv(x) + + return x + + +class EqualLinear(nn.Module): + def __init__(self, dim, dim_out, lr_mul=1, bias=True): + super().__init__() + self.weight = nn.Parameter(torch.randn(dim_out, dim)) + if bias: + self.bias = nn.Parameter(torch.zeros(dim_out)) + + self.lr_mul = lr_mul + + def forward(self, input): + return F.linear(input, self.weight * self.lr_mul, bias=self.bias * self.lr_mul) + + +class StyleGanNetwork(nn.Module): + def __init__(self, dim_in=128, dim_out=512, depth=8, lr_mul=0.1, dim_text_latent=0): + super().__init__() + self.dim_in = dim_in + self.dim_out = dim_out + self.dim_text_latent = dim_text_latent + + layers = [] + for i in range(depth): + is_first = i == 0 + + if is_first: + dim_in_layer = dim_in + dim_text_latent + else: + dim_in_layer = dim_out + + dim_out_layer = dim_out + + layers.extend( + [EqualLinear(dim_in_layer, dim_out_layer, lr_mul), nn.LeakyReLU(0.2)] + ) + + self.net = nn.Sequential(*layers) + + def forward(self, x, text_latent=None): + x = F.normalize(x, dim=1) + if self.dim_text_latent > 0: + assert exists(text_latent) + x = torch.cat((x, text_latent), dim=-1) + return self.net(x) + + +class UnetUpsampler(torch.nn.Module): + + def __init__( + self, + dim: int, + *, + image_size: int, + input_image_size: int, + init_dim: Optional[int] = None, + out_dim: Optional[int] = None, + style_network: Optional[dict] = None, + up_dim_mults: tuple = (1, 2, 4, 8, 16), + down_dim_mults: tuple = (4, 8, 16), + channels: int = 3, + resnet_block_groups: int = 8, + full_attn: tuple = (False, False, False, True, True), + flash_attn: bool = True, + self_attn_dim_head: int = 64, + self_attn_heads: int = 8, + attn_depths: tuple = (2, 2, 2, 2, 4), + mid_attn_depth: int = 4, + num_conv_kernels: int = 4, + resize_mode: str = "bilinear", + unconditional: bool = True, + skip_connect_scale: Optional[float] = None, + ): + super().__init__() + self.style_network = style_network = StyleGanNetwork(**style_network) + self.unconditional = unconditional + assert not ( + unconditional + and exists(style_network) + and style_network.dim_text_latent > 0 + ) + + assert is_power_of_two(image_size) and is_power_of_two( + input_image_size + ), "both output image size and input image size must be power of 2" + assert ( + input_image_size < image_size + ), "input image size must be smaller than the output image size, thus upsampling" + + self.image_size = image_size + self.input_image_size = input_image_size + + style_embed_split_dims = [] + + self.channels = channels + input_channels = channels + + init_dim = default(init_dim, dim) + + up_dims = [init_dim, *map(lambda m: dim * m, up_dim_mults)] + init_down_dim = up_dims[len(up_dim_mults) - len(down_dim_mults)] + down_dims = [init_down_dim, *map(lambda m: dim * m, down_dim_mults)] + self.init_conv = nn.Conv2d(input_channels, init_down_dim, 7, padding=3) + + up_in_out = list(zip(up_dims[:-1], up_dims[1:])) + down_in_out = list(zip(down_dims[:-1], down_dims[1:])) + + block_klass = partial( + ResnetBlock, + groups=resnet_block_groups, + num_conv_kernels=num_conv_kernels, + style_dims=style_embed_split_dims, + ) + + FullAttention = partial(Transformer, flash_attn=flash_attn) + *_, mid_dim = up_dims + + self.skip_connect_scale = default(skip_connect_scale, 2 ** -0.5) + + self.downs = nn.ModuleList([]) + self.ups = nn.ModuleList([]) + + block_count = 6 + + for ind, ( + (dim_in, dim_out), + layer_full_attn, + layer_attn_depth, + ) in enumerate(zip(down_in_out, full_attn, attn_depths)): + attn_klass = FullAttention if layer_full_attn else LinearTransformer + + blocks = [] + for i in range(block_count): + blocks.append(block_klass(dim_in, dim_in)) + + self.downs.append( + nn.ModuleList( + [ + nn.ModuleList(blocks), + nn.ModuleList( + [ + ( + attn_klass( + dim_in, + dim_head=self_attn_dim_head, + heads=self_attn_heads, + depth=layer_attn_depth, + ) + if layer_full_attn + else None + ), + nn.Conv2d( + dim_in, dim_out, kernel_size=3, stride=2, padding=1 + ), + ] + ), + ] + ) + ) + + self.mid_block1 = block_klass(mid_dim, mid_dim) + self.mid_attn = FullAttention( + mid_dim, + dim_head=self_attn_dim_head, + heads=self_attn_heads, + depth=mid_attn_depth, + ) + self.mid_block2 = block_klass(mid_dim, mid_dim) + + *_, last_dim = up_dims + + for ind, ( + (dim_in, dim_out), + layer_full_attn, + layer_attn_depth, + ) in enumerate( + zip( + reversed(up_in_out), + reversed(full_attn), + reversed(attn_depths), + ) + ): + attn_klass = FullAttention if layer_full_attn else LinearTransformer + + blocks = [] + input_dim = dim_in * 2 if ind < len(down_in_out) else dim_in + for i in range(block_count): + blocks.append(block_klass(input_dim, dim_in)) + + self.ups.append( + nn.ModuleList( + [ + nn.ModuleList(blocks), + nn.ModuleList( + [ + NearestNeighborhoodUpsample( + last_dim if ind == 0 else dim_out, + dim_in, + ), + ( + attn_klass( + dim_in, + dim_head=self_attn_dim_head, + heads=self_attn_heads, + depth=layer_attn_depth, + ) + if layer_full_attn + else None + ), + ] + ), + ] + ) + ) + + self.out_dim = default(out_dim, channels) + self.final_res_block = block_klass(dim, dim) + self.final_to_rgb = nn.Conv2d(dim, channels, 1) + self.resize_mode = resize_mode + self.style_to_conv_modulations = nn.Linear( + style_network.dim_out, sum(style_embed_split_dims) + ) + self.style_embed_split_dims = style_embed_split_dims + + @property + def allowable_rgb_resolutions(self): + input_res_base = int(log2(self.input_image_size)) + output_res_base = int(log2(self.image_size)) + allowed_rgb_res_base = list(range(input_res_base, output_res_base)) + return [*map(lambda p: 2 ** p, allowed_rgb_res_base)] + + @property + def device(self): + return next(self.parameters()).device + + @property + def total_params(self): + return sum([p.numel() for p in self.parameters()]) + + def resize_image_to(self, x, size): + return F.interpolate(x, (size, size), mode=self.resize_mode) + + def forward( + self, + lowres_image: torch.Tensor, + styles: Optional[torch.Tensor] = None, + noise: Optional[torch.Tensor] = None, + global_text_tokens: Optional[torch.Tensor] = None, + return_all_rgbs: bool = False, + ): + x = lowres_image + + noise_scale = 0.001 # Adjust the scale of the noise as needed + noise_aug = torch.randn_like(x) * noise_scale + x = x + noise_aug + x = x.clamp(0, 1) + + shape = x.shape + batch_size = shape[0] + + assert shape[-2:] == ((self.input_image_size,) * 2) + + # styles + if not exists(styles): + assert exists(self.style_network) + + noise = default( + noise, + torch.randn( + (batch_size, self.style_network.dim_in), device=self.device + ), + ) + styles = self.style_network(noise, global_text_tokens) + + # project styles to conv modulations + conv_mods = self.style_to_conv_modulations(styles) + conv_mods = conv_mods.split(self.style_embed_split_dims, dim=-1) + conv_mods = iter(conv_mods) + + x = self.init_conv(x) + + h = [] + for blocks, (attn, downsample) in self.downs: + for block in blocks: + x = block(x, conv_mods_iter=conv_mods) + h.append(x) + + if attn is not None: + x = attn(x) + + x = downsample(x) + + x = self.mid_block1(x, conv_mods_iter=conv_mods) + x = self.mid_attn(x) + x = self.mid_block2(x, conv_mods_iter=conv_mods) + + for ( + blocks, + ( + upsample, + attn, + ), + ) in self.ups: + x = upsample(x) + for block in blocks: + if h != []: + res = h.pop() + res = res * self.skip_connect_scale + x = torch.cat((x, res), dim=1) + + x = block(x, conv_mods_iter=conv_mods) + + if attn is not None: + x = attn(x) + + x = self.final_res_block(x, conv_mods_iter=conv_mods) + rgb = self.final_to_rgb(x) + + if not return_all_rgbs: + return rgb + + return rgb, [] + + +def tile_image(image, chunk_size=64): + c, h, w = image.shape + h_chunks = ceil(h / chunk_size) + w_chunks = ceil(w / chunk_size) + tiles = [] + for i in range(h_chunks): + for j in range(w_chunks): + tile = image[:, i * chunk_size:(i + 1) * chunk_size, j * chunk_size:(j + 1) * chunk_size] + tiles.append(tile) + return tiles, h_chunks, w_chunks + + +# This helps create a checkboard pattern with some edge blending +def create_checkerboard_weights(input_image_size): + x = torch.linspace(-1, 1, input_image_size) + y = torch.linspace(-1, 1, input_image_size) + + x, y = torch.meshgrid(x, y, indexing='ij') + d = torch.sqrt(x * x + y * y) + sigma, mu = 0.5, 0.0 + weights = torch.exp(-((d - mu) ** 2 / (2.0 * sigma ** 2))) + + # saturate the values to sure get high weights in the center + weights = weights ** 8 + + return weights / weights.max() # Normalize to [0, 1] + + +def repeat_weights(weights, image_size): + input_image_size = weights.shape[0] + repeats = (math.ceil(image_size[0] / input_image_size), math.ceil(image_size[1] / input_image_size)) + return weights.repeat(repeats)[:image_size[0], :image_size[1]] + + +def create_offset_weights(weights, image_size): + input_image_size = weights.shape[0] + offset = input_image_size // 2 + full_weights = repeat_weights(weights, (image_size[0] + offset, image_size[1] + offset)) + return full_weights[offset:, offset:] + + +def merge_tiles(tiles, h_chunks, w_chunks, chunk_size=64): + # Determine the shape of the output tensor + c = tiles[0].shape[0] + h = h_chunks * chunk_size + w = w_chunks * chunk_size + + # Create an empty tensor to hold the merged image + merged = torch.zeros((c, h, w), dtype=tiles[0].dtype) + + # Iterate over the tiles and place them in the correct position + for idx, tile in enumerate(tiles): + i = idx // w_chunks + j = idx % w_chunks + + h_start = i * chunk_size + w_start = j * chunk_size + + tile_h, tile_w = tile.shape[1:] + merged[:, h_start:h_start + tile_h, w_start:w_start + tile_w] = tile + + return merged + +class Upscaler(nn.Module): + def __init__(self): + super(Upscaler, self).__init__() + + self.name = 'upscaler' + self.config = { + "style_network": { + "dim_in": 128, + "dim_out": 512, + "depth": 4 + }, + "dim": 64, + "image_size": 256, + "input_image_size": 64, + "unconditional": True, + "skip_connect_scale": 0.4 + } + + self.generator = UnetUpsampler(**self.config) + + def forward(self, x, y): + return self.generator(x, noise=y) + + @property + def device(self): + return self.generator.device + + +@torch.no_grad() +def upscale_4x(image, upsampler, input_image_size=64, max_batch_size=8): + tensor_transform = transforms.ToTensor() + device = upsampler.device + + image_tensor = tensor_transform(image).unsqueeze(0) + _, _, h, w = image_tensor.shape + pad_h = (input_image_size - h % input_image_size) % input_image_size + pad_w = (input_image_size - w % input_image_size) % input_image_size + + # Pad the image + image_tensor = torch.nn.functional.pad(image_tensor, (0, pad_w, 0, pad_h), mode='reflect').squeeze(0) + tiles, h_chunks, w_chunks = tile_image(image_tensor, input_image_size) + + # Batch processing of tiles + num_tiles = len(tiles) + batches = [tiles[i:i + max_batch_size] for i in range(0, num_tiles, max_batch_size)] + reconstructed_tiles = [] + + for batch in batches: + model_input = torch.stack(batch).to(device) + generator_output = upsampler( + model_input, + torch.randn(model_input.shape[0], 128, device=device) + ) + reconstructed_tiles.extend(list(generator_output.clamp_(0, 1).detach().cpu())) + + merged_tensor = merge_tiles(reconstructed_tiles, h_chunks, w_chunks, input_image_size * 4) + unpadded = merged_tensor[:, :h * 4, :w * 4] + + to_pil = transforms.ToPILImage() + return np.array(to_pil(unpadded)) + +# Tiled 4x upscaling with overlapping tiles to reduce seam artifacts +# weights options are 'checkboard' and 'constant' +@torch.no_grad() +def upscale_4x_overlapped(image, upsampler, input_image_size=64, max_batch_size=8, weight_type='checkboard'): + tensor_transform = transforms.ToTensor() + device = upsampler.device + + image_tensor = tensor_transform(image).unsqueeze(0) + _, _, h, w = image_tensor.shape + + # Calculate paddings + pad_h = ( + input_image_size - h % input_image_size + ) % input_image_size + pad_w = ( + input_image_size - w % input_image_size + ) % input_image_size + + # Pad the image + image_tensor = torch.nn.functional.pad( + image_tensor, (0, pad_w, 0, pad_h), mode="reflect" + ).squeeze(0) + + # Function to process tiles + def process_tiles(tiles, h_chunks, w_chunks): + num_tiles = len(tiles) + batches = [ + tiles[i: i + max_batch_size] + for i in range(0, num_tiles, max_batch_size) + ] + reconstructed_tiles = [] + + for batch in batches: + model_input = torch.stack(batch).to(device) + generator_output = upsampler(model_input, + torch.randn(model_input.shape[0], 128, device=device)) + reconstructed_tiles.extend(list(generator_output.clamp_(0, 1).detach().cpu())) + + return merge_tiles( + reconstructed_tiles, h_chunks, w_chunks, input_image_size * 4 + ) + + # First pass + tiles1, h_chunks1, w_chunks1 = tile_image(image_tensor, input_image_size) + result1 = process_tiles(tiles1, h_chunks1, w_chunks1) + + # Second pass with offset + offset = input_image_size // 2 + image_tensor_offset = torch.nn.functional.pad(image_tensor, (offset, offset, offset, offset), + mode='reflect').squeeze(0) + + tiles2, h_chunks2, w_chunks2 = tile_image( + image_tensor_offset, input_image_size + ) + result2 = process_tiles(tiles2, h_chunks2, w_chunks2) + + # unpad + offset_4x = offset * 4 + result2_interior = result2[:, offset_4x:-offset_4x, offset_4x:-offset_4x] + + if weight_type == 'checkboard': + weight_tile = create_checkerboard_weights(input_image_size * 4) + + weight_shape = result2_interior.shape[1:] + weights_1 = create_offset_weights(weight_tile, weight_shape) + weights_2 = repeat_weights(weight_tile, weight_shape) + + normalizer = weights_1 + weights_2 + weights_1 = weights_1 / normalizer + weights_2 = weights_2 / normalizer + + weights_1 = weights_1.unsqueeze(0).repeat(3, 1, 1) + weights_2 = weights_2.unsqueeze(0).repeat(3, 1, 1) + elif weight_type == 'constant': + weights_1 = torch.ones_like(result2_interior) * 0.5 + weights_2 = weights_1 + else: + raise ValueError("weight_type should be either 'gaussian' or 'constant' but got", weight_type) + + result1 = result1 * weights_2 + result2 = result2_interior * weights_1 + + # Average the overlapping region + result1 = ( + result1 + result2 + ) + + # Remove padding + unpadded = result1[:, : h * 4, : w * 4] + + to_pil = transforms.ToPILImage() + return np.array(to_pil(unpadded)) diff --git a/networks/colorizer.py b/networks/colorizer.py new file mode 100644 index 0000000..fed19f1 --- /dev/null +++ b/networks/colorizer.py @@ -0,0 +1,317 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.nn import Parameter + +from .extractor import SEResNeXt_Origin, BottleneckX_Origin + +'''https://github.com/orashi/AlacGAN/blob/master/models/standard.py''' + +def l2normalize(v, eps=1e-12): + return v / (v.norm() + eps) + + +class SpectralNorm(nn.Module): + def __init__(self, module, name='weight', power_iterations=1): + super(SpectralNorm, self).__init__() + self.module = module + self.name = name + self.power_iterations = power_iterations + if not self._made_params(): + self._make_params() + + def _update_u_v(self): + u = getattr(self.module, self.name + "_u") + v = getattr(self.module, self.name + "_v") + w = getattr(self.module, self.name + "_bar") + + height = w.data.shape[0] + for _ in range(self.power_iterations): + v.data = l2normalize(torch.mv(torch.t(w.view(height,-1).data), u.data)) + u.data = l2normalize(torch.mv(w.view(height,-1).data, v.data)) + + # sigma = torch.dot(u.data, torch.mv(w.view(height,-1).data, v.data)) + sigma = u.dot(w.view(height, -1).mv(v)) + setattr(self.module, self.name, w / sigma.expand_as(w)) + + def _made_params(self): + try: + u = getattr(self.module, self.name + "_u") + v = getattr(self.module, self.name + "_v") + w = getattr(self.module, self.name + "_bar") + return True + except AttributeError: + return False + + + def _make_params(self): + w = getattr(self.module, self.name) + height = w.data.shape[0] + width = w.view(height, -1).data.shape[1] + + u = Parameter(w.data.new(height).normal_(0, 1), requires_grad=False) + v = Parameter(w.data.new(width).normal_(0, 1), requires_grad=False) + u.data = l2normalize(u.data) + v.data = l2normalize(v.data) + w_bar = Parameter(w.data) + + del self.module._parameters[self.name] + + self.module.register_parameter(self.name + "_u", u) + self.module.register_parameter(self.name + "_v", v) + self.module.register_parameter(self.name + "_bar", w_bar) + + + def forward(self, *args): + self._update_u_v() + return self.module.forward(*args) + +class Selayer(nn.Module): + def __init__(self, inplanes): + super(Selayer, self).__init__() + self.global_avgpool = nn.AdaptiveAvgPool2d(1) + self.conv1 = nn.Conv2d(inplanes, inplanes // 16, kernel_size=1, stride=1) + self.conv2 = nn.Conv2d(inplanes // 16, inplanes, kernel_size=1, stride=1) + self.relu = nn.ReLU(inplace=True) + self.sigmoid = nn.Sigmoid() + + def forward(self, x): + out = self.global_avgpool(x) + out = self.conv1(out) + out = self.relu(out) + out = self.conv2(out) + out = self.sigmoid(out) + + return x * out + +class SelayerSpectr(nn.Module): + def __init__(self, inplanes): + super(SelayerSpectr, self).__init__() + self.global_avgpool = nn.AdaptiveAvgPool2d(1) + self.conv1 = SpectralNorm(nn.Conv2d(inplanes, inplanes // 16, kernel_size=1, stride=1)) + self.conv2 = SpectralNorm(nn.Conv2d(inplanes // 16, inplanes, kernel_size=1, stride=1)) + self.relu = nn.ReLU(inplace=True) + self.sigmoid = nn.Sigmoid() + + def forward(self, x): + out = self.global_avgpool(x) + out = self.conv1(out) + out = self.relu(out) + out = self.conv2(out) + out = self.sigmoid(out) + + return x * out + +class ResNeXtBottleneck(nn.Module): + def __init__(self, in_channels=256, out_channels=256, stride=1, cardinality=32, dilate=1): + super(ResNeXtBottleneck, self).__init__() + D = out_channels // 2 + self.out_channels = out_channels + self.conv_reduce = nn.Conv2d(in_channels, D, kernel_size=1, stride=1, padding=0, bias=False) + self.conv_conv = nn.Conv2d(D, D, kernel_size=2 + stride, stride=stride, padding=dilate, dilation=dilate, + groups=cardinality, + bias=False) + self.conv_expand = nn.Conv2d(D, out_channels, kernel_size=1, stride=1, padding=0, bias=False) + self.shortcut = nn.Sequential() + if stride != 1: + self.shortcut.add_module('shortcut', + nn.AvgPool2d(2, stride=2)) + + self.selayer = Selayer(out_channels) + + def forward(self, x): + bottleneck = self.conv_reduce.forward(x) + bottleneck = F.leaky_relu(bottleneck, 0.2, True) + bottleneck = self.conv_conv.forward(bottleneck) + bottleneck = F.leaky_relu(bottleneck, 0.2, True) + bottleneck = self.conv_expand.forward(bottleneck) + bottleneck = self.selayer(bottleneck) + + x = self.shortcut.forward(x) + return x + bottleneck + +class SpectrResNeXtBottleneck(nn.Module): + def __init__(self, in_channels=256, out_channels=256, stride=1, cardinality=32, dilate=1): + super(SpectrResNeXtBottleneck, self).__init__() + D = out_channels // 2 + self.out_channels = out_channels + self.conv_reduce = SpectralNorm(nn.Conv2d(in_channels, D, kernel_size=1, stride=1, padding=0, bias=False)) + self.conv_conv = SpectralNorm(nn.Conv2d(D, D, kernel_size=2 + stride, stride=stride, padding=dilate, dilation=dilate, + groups=cardinality, + bias=False)) + self.conv_expand = SpectralNorm(nn.Conv2d(D, out_channels, kernel_size=1, stride=1, padding=0, bias=False)) + self.shortcut = nn.Sequential() + if stride != 1: + self.shortcut.add_module('shortcut', + nn.AvgPool2d(2, stride=2)) + + self.selayer = SelayerSpectr(out_channels) + + def forward(self, x): + bottleneck = self.conv_reduce.forward(x) + bottleneck = F.leaky_relu(bottleneck, 0.2, True) + bottleneck = self.conv_conv.forward(bottleneck) + bottleneck = F.leaky_relu(bottleneck, 0.2, True) + bottleneck = self.conv_expand.forward(bottleneck) + bottleneck = self.selayer(bottleneck) + + x = self.shortcut.forward(x) + return x + bottleneck + +class FeatureConv(nn.Module): + def __init__(self, input_dim=512, output_dim=512): + super(FeatureConv, self).__init__() + + no_bn = True + + seq = [] + seq.append(nn.Conv2d(input_dim, output_dim, kernel_size=3, stride=1, padding=1, bias=False)) + if not no_bn: seq.append(nn.BatchNorm2d(output_dim)) + seq.append(nn.ReLU(inplace=True)) + seq.append(nn.Conv2d(output_dim, output_dim, kernel_size=3, stride=2, padding=1, bias=False)) + if not no_bn: seq.append(nn.BatchNorm2d(output_dim)) + seq.append(nn.ReLU(inplace=True)) + seq.append(nn.Conv2d(output_dim, output_dim, kernel_size=3, stride=1, padding=1, bias=False)) + seq.append(nn.ReLU(inplace=True)) + + self.network = nn.Sequential(*seq) + + def forward(self, x): + return self.network(x) + +class Generator(nn.Module): + def __init__(self, ngf=64): + super(Generator, self).__init__() + + self.encoder = SEResNeXt_Origin(BottleneckX_Origin, [3, 4, 6, 3], num_classes= 370, input_channels=1) + + self.to0 = self._make_encoder_block_first(5, 32) + self.to1 = self._make_encoder_block(32, 64) + self.to2 = self._make_encoder_block(64, 92) + self.to3 = self._make_encoder_block(92, 128) + self.to4 = self._make_encoder_block(128, 256) + + self.deconv_for_decoder = nn.Sequential( + nn.ConvTranspose2d(256, 128, 3, stride=2, padding=1, output_padding=1), # output is 64 * 64 + nn.LeakyReLU(0.2), + nn.ConvTranspose2d(128, 64, 3, stride=2, padding=1, output_padding=1), # output is 128 * 128 + nn.LeakyReLU(0.2), + nn.ConvTranspose2d(64, 32, 3, stride=1, padding=1, output_padding=0), # output is 256 * 256 + nn.LeakyReLU(0.2), + nn.ConvTranspose2d(32, 3, 3, stride=1, padding=1, output_padding=0), # output is 256 * 256 + nn.Tanh(), + ) + + tunnel4 = nn.Sequential(*[ResNeXtBottleneck(512, 512, cardinality=32, dilate=1) for _ in range(20)]) + + + self.tunnel4 = nn.Sequential(nn.Conv2d(1024 + 128, 512, kernel_size=3, stride=1, padding=1), + nn.LeakyReLU(0.2, True), + tunnel4, + nn.Conv2d(512, 1024, kernel_size=3, stride=1, padding=1), + nn.PixelShuffle(2), + nn.LeakyReLU(0.2, True) + ) # 64 + + depth = 2 + tunnel = [ResNeXtBottleneck(256, 256, cardinality=32, dilate=1) for _ in range(depth)] + tunnel += [ResNeXtBottleneck(256, 256, cardinality=32, dilate=2) for _ in range(depth)] + tunnel += [ResNeXtBottleneck(256, 256, cardinality=32, dilate=4) for _ in range(depth)] + tunnel += [ResNeXtBottleneck(256, 256, cardinality=32, dilate=2), + ResNeXtBottleneck(256, 256, cardinality=32, dilate=1)] + tunnel3 = nn.Sequential(*tunnel) + + self.tunnel3 = nn.Sequential(nn.Conv2d(512 + 256, 256, kernel_size=3, stride=1, padding=1), + nn.LeakyReLU(0.2, True), + tunnel3, + nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1), + nn.PixelShuffle(2), + nn.LeakyReLU(0.2, True) + ) # 128 + + tunnel = [ResNeXtBottleneck(128, 128, cardinality=32, dilate=1) for _ in range(depth)] + tunnel += [ResNeXtBottleneck(128, 128, cardinality=32, dilate=2) for _ in range(depth)] + tunnel += [ResNeXtBottleneck(128, 128, cardinality=32, dilate=4) for _ in range(depth)] + tunnel += [ResNeXtBottleneck(128, 128, cardinality=32, dilate=2), + ResNeXtBottleneck(128, 128, cardinality=32, dilate=1)] + tunnel2 = nn.Sequential(*tunnel) + + self.tunnel2 = nn.Sequential(nn.Conv2d(128 + 256 + 64, 128, kernel_size=3, stride=1, padding=1), + nn.LeakyReLU(0.2, True), + tunnel2, + nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1), + nn.PixelShuffle(2), + nn.LeakyReLU(0.2, True) + ) + + tunnel = [ResNeXtBottleneck(64, 64, cardinality=16, dilate=1)] + tunnel += [ResNeXtBottleneck(64, 64, cardinality=16, dilate=2)] + tunnel += [ResNeXtBottleneck(64, 64, cardinality=16, dilate=4)] + tunnel += [ResNeXtBottleneck(64, 64, cardinality=16, dilate=2), + ResNeXtBottleneck(64, 64, cardinality=16, dilate=1)] + tunnel1 = nn.Sequential(*tunnel) + + self.tunnel1 = nn.Sequential(nn.Conv2d(64 + 32, 64, kernel_size=3, stride=1, padding=1), + nn.LeakyReLU(0.2, True), + tunnel1, + nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1), + nn.PixelShuffle(2), + nn.LeakyReLU(0.2, True) + ) + + self.exit = nn.Sequential(nn.Conv2d(64 + 32, 32, kernel_size=3, stride=1, padding=1), + nn.LeakyReLU(0.2, True), + nn.Conv2d(32, 3, kernel_size= 1, stride = 1, padding = 0)) + + + def _make_encoder_block(self, inplanes, planes): + return nn.Sequential( + nn.Conv2d(inplanes, planes, 3, 2, 1), + nn.LeakyReLU(0.2), + nn.Conv2d(planes, planes, 3, 1, 1), + nn.LeakyReLU(0.2), + ) + + def _make_encoder_block_first(self, inplanes, planes): + return nn.Sequential( + nn.Conv2d(inplanes, planes, 3, 1, 1), + nn.LeakyReLU(0.2), + nn.Conv2d(planes, planes, 3, 1, 1), + nn.LeakyReLU(0.2), + ) + + def forward(self, sketch): + + x0 = self.to0(sketch) + aux_out = self.to1(x0) + aux_out = self.to2(aux_out) + aux_out = self.to3(aux_out) + + x1, x2, x3, x4 = self.encoder(sketch[:, 0:1]) + + out = self.tunnel4(torch.cat([x4, aux_out], 1)) + + + + x = self.tunnel3(torch.cat([out, x3], 1)) + + x = self.tunnel2(torch.cat([x, x2, x1], 1)) + + + x = torch.tanh(self.exit(torch.cat([x, x0], 1))) + + decoder_output = self.deconv_for_decoder(out) + + return x, decoder_output + + +class Colorizer(nn.Module): + def __init__(self): + super(Colorizer, self).__init__() + + self.name = 'colorizer' + self.generator = Generator() + + def forward(self, x, extractor_grad = False): + fake, guide = self.generator(x) + return fake, guide diff --git a/networks/extractor.py b/networks/extractor.py new file mode 100644 index 0000000..e7e6ec5 --- /dev/null +++ b/networks/extractor.py @@ -0,0 +1,126 @@ +import torch.nn as nn +import math + +'''https://github.com/blandocs/Tag2Pix/blob/master/model/pretrained.py''' + +# Pretrained version +class Selayer(nn.Module): + def __init__(self, inplanes): + super(Selayer, self).__init__() + self.global_avgpool = nn.AdaptiveAvgPool2d(1) + self.conv1 = nn.Conv2d(inplanes, inplanes // 16, kernel_size=1, stride=1) + self.conv2 = nn.Conv2d(inplanes // 16, inplanes, kernel_size=1, stride=1) + self.relu = nn.ReLU(inplace=True) + self.sigmoid = nn.Sigmoid() + + def forward(self, x): + out = self.global_avgpool(x) + out = self.conv1(out) + out = self.relu(out) + out = self.conv2(out) + out = self.sigmoid(out) + + return x * out + + +class BottleneckX_Origin(nn.Module): + expansion = 4 + + def __init__(self, inplanes, planes, cardinality, stride=1, downsample=None): + super(BottleneckX_Origin, self).__init__() + self.conv1 = nn.Conv2d(inplanes, planes * 2, kernel_size=1, bias=False) + self.bn1 = nn.BatchNorm2d(planes * 2) + + self.conv2 = nn.Conv2d(planes * 2, planes * 2, kernel_size=3, stride=stride, + padding=1, groups=cardinality, bias=False) + self.bn2 = nn.BatchNorm2d(planes * 2) + + self.conv3 = nn.Conv2d(planes * 2, planes * 4, kernel_size=1, bias=False) + self.bn3 = nn.BatchNorm2d(planes * 4) + + self.selayer = Selayer(planes * 4) + + self.relu = nn.ReLU(inplace=True) + self.downsample = downsample + self.stride = stride + + def forward(self, x): + residual = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + + out = self.conv2(out) + out = self.bn2(out) + out = self.relu(out) + + out = self.conv3(out) + out = self.bn3(out) + + out = self.selayer(out) + + if self.downsample is not None: + residual = self.downsample(x) + + out += residual + out = self.relu(out) + + return out + +class SEResNeXt_Origin(nn.Module): + def __init__(self, block, layers, input_channels=3, cardinality=32, num_classes=1000): + super(SEResNeXt_Origin, self).__init__() + self.cardinality = cardinality + self.inplanes = 64 + self.input_channels = input_channels + + self.conv1 = nn.Conv2d(input_channels, 64, kernel_size=7, stride=2, padding=3, + bias=False) + self.bn1 = nn.BatchNorm2d(64) + self.relu = nn.ReLU(inplace=True) + + self.layer1 = self._make_layer(block, 64, layers[0]) + self.layer2 = self._make_layer(block, 128, layers[1], stride=2) + self.layer3 = self._make_layer(block, 256, layers[2], stride=2) + + for m in self.modules(): + if isinstance(m, nn.Conv2d): + n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + m.weight.data.normal_(0, math.sqrt(2. / n)) + if m.bias is not None: + m.bias.data.zero_() + elif isinstance(m, nn.BatchNorm2d): + m.weight.data.fill_(1) + m.bias.data.zero_() + + def _make_layer(self, block, planes, blocks, stride=1): + downsample = None + if stride != 1 or self.inplanes != planes * block.expansion: + downsample = nn.Sequential( + nn.Conv2d(self.inplanes, planes * block.expansion, + kernel_size=1, stride=stride, bias=False), + nn.BatchNorm2d(planes * block.expansion), + ) + + layers = [] + layers.append(block(self.inplanes, planes, self.cardinality, stride, downsample)) + self.inplanes = planes * block.expansion + for i in range(1, blocks): + layers.append(block(self.inplanes, planes, self.cardinality)) + + return nn.Sequential(*layers) + + def forward(self, x): + + x = self.conv1(x) + x = self.bn1(x) + x1 = self.relu(x) + + x2 = self.layer1(x1) + + x3 = self.layer2(x2) + + x4 = self.layer3(x3) + + return x1, x2, x3, x4 diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..53c677b --- /dev/null +++ b/requirements.txt @@ -0,0 +1,14 @@ +# Core Processing & ML Libraries +numpy +Pillow +opencv-python +scikit-image +einops +matplotlib +mangadex-downloader +nhentai + +# GUI and Theming +sv-ttk +darkdetect +pywinstyles; sys_platform == 'win32' diff --git a/train.bat b/train.bat new file mode 100644 index 0000000..0e8884c --- /dev/null +++ b/train.bat @@ -0,0 +1,28 @@ +@echo off +title Manga Colorizer Training +echo ======================================= +echo Manga Colorizer - Training Launcher +echo ======================================= + +cd /d "C:\Users\Nighthawk\Desktop\manga_colorize" + +echo. +echo [1/1] Activating virtual environment... +call .venv\Scripts\activate +if %errorlevel% neq 0 ( + echo ERROR: Failed to activate virtual environment. + pause + exit /b +) + +echo. +echo ======================================= +echo Launching Training Script... +echo ======================================= +echo. + +python train_model.py + +echo. +echo Training script has exited. Press any key to close this window. +pause >nul diff --git a/train_model.py b/train_model.py new file mode 100644 index 0000000..e9a2f06 --- /dev/null +++ b/train_model.py @@ -0,0 +1,593 @@ +# train_model.py — Rich UI fine-tuner (Windows-safe, CSV+plot logging, graceful) + +import os, sys, time, json, random, platform, logging, csv, datetime, shutil +from pathlib import Path +from typing import Optional + +# ---- Paths (edit if needed) ---- +GEN_ZIP = r"C:\Users\Nighthawk\Desktop\manga_colorize\networks\generator.zip" +OUT_DIR = r"C:\Users\Nighthawk\Desktop\manga_colorize\finetune_out" +PAUSE_ON_EXIT = False # set True if you want a final "Press Enter..." pause + +IS_WINDOWS = platform.system() == "Windows" + +# ---- Third-party ---- +import numpy as np +from PIL import Image, ImageFilter, ImageDraw + +import torch +import torch.nn as nn +from torch.utils.data import Dataset, DataLoader +from torch import amp as torch_amp # modern AMP API + +# Optional plotting +HAVE_MPL = True +try: + import matplotlib.pyplot as plt +except Exception: + HAVE_MPL = False + +# HuggingFace datasets (optional) +HAVE_DATASETS = True +try: + from datasets import load_dataset +except Exception: + HAVE_DATASETS = False + +# Rich UI +from rich.console import Console +from rich.panel import Panel +from rich.table import Table +from rich.prompt import Prompt, Confirm +from rich.progress import ( + Progress, TextColumn, BarColumn, TimeElapsedColumn, + TimeRemainingColumn +) +from rich.logging import RichHandler + +console = Console() + +# Import your model from the repo +sys.path.insert(0, str(Path(__file__).parent)) +from networks.colorizer import Colorizer + + +# ----------------- Logging ----------------- + +def setup_logging(out_dir: str): + Path(out_dir).mkdir(parents=True, exist_ok=True) + log_path = Path(out_dir) / "train.log" + + logger = logging.getLogger() + logger.setLevel(logging.INFO) + for h in list(logger.handlers): + logger.removeHandler(h) + + ch = RichHandler(console=console, show_path=False, rich_tracebacks=True) + ch.setLevel(logging.INFO) + ch.setFormatter(logging.Formatter("%(message)s")) + + fh = logging.FileHandler(log_path, encoding="utf-8") + fh.setLevel(logging.INFO) + fh.setFormatter(logging.Formatter( + "%(asctime)s | %(levelname)s | %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + )) + + logger.addHandler(ch) + logger.addHandler(fh) + logging.info(f"Logging to: {log_path}") + + +# ----------------- Utilities ----------------- + +def ensure_dir(p: str): + Path(p).mkdir(parents=True, exist_ok=True) + +def set_seed(seed: int = 42, deterministic: bool = True): + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + if deterministic: + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + +def device_info(device: str): + if device == "cuda": + name = torch.cuda.get_device_name(0) + cap = torch.cuda.get_device_capability(0) + mem = torch.cuda.get_device_properties(0).total_memory / (1024**3) + return f"[bold cyan]CUDA[/]: {name} (CC {cap[0]}.{cap[1]}), {mem:.1f} GB" + return "[bold yellow]CPU[/]" + +def mem_stats(): + if torch.cuda.is_available(): + alloc = torch.cuda.memory_allocated() / (1024**2) + reserv = torch.cuda.memory_reserved() / (1024**2) + return f"VRAM alloc {alloc:.0f}MB / reserved {reserv:.0f}MB" + return "CPU memory" + +def to_uint8(img01: np.ndarray) -> np.ndarray: + return np.clip(img01 * 255.0, 0, 255).round().astype(np.uint8) + +def save_image_row(path: str, images: list): + if not images: return + h, w, _ = images[0].shape + from PIL import Image as _I + grid = _I.new("RGB", (w * len(images), h)) + x = 0 + for im in images: + grid.paste(_I.fromarray(im), (x, 0)) + x += w + grid.save(path) + +def rgb_to_L(rgb_np: np.ndarray) -> np.ndarray: + return np.expand_dims(np.dot(rgb_np[..., :3], [0.299, 0.587, 0.114]), 2) + +def random_scribbles(w, h, n=3): + hint = Image.new("RGB", (w, h), (128, 128, 128)) + mask = Image.new("L", (w, h), 0) + dh, dm = ImageDraw.Draw(hint), ImageDraw.Draw(mask) + for _ in range(n): + color = tuple(np.random.randint(0, 256, 3).tolist()) + import numpy as _np + pts = [(_np.random.randint(0, w), _np.random.randint(0, h))] + for _ in range(8): + x = int(np.clip(pts[-1][0] + _np.random.randint(-w // 6, w // 6), 0, w - 1)) + y = int(np.clip(pts[-1][1] + _np.random.randint(-h // 6, h // 6), 0, h - 1)) + pts.append((x, y)) + width = _np.random.randint(8, 20) + dh.line(pts, fill=color, width=width) + dm.line(pts, fill=255, width=width) + hint = hint.filter(ImageFilter.GaussianBlur(1.0)) + return np.asarray(hint), np.asarray(mask) + +def state_dict_load(path: str): + logging.info(f"Loading weights: {path}") + sd = torch.load(path, map_location="cpu") + if not isinstance(sd, dict): + raise RuntimeError("Loaded object is not a state_dict (dict).") + return sd + +def latest_ckpt(dir_path: str) -> Optional[Path]: + d = Path(dir_path) + if not d.exists(): return None + cks = sorted(d.glob("ckpt_step*.pt"), key=lambda p: p.stat().st_mtime) + return cks[-1] if cks else None + +def rotate_checkpoints(dir_path: str, keep: int = 5): + d = Path(dir_path) + cks = sorted(d.glob("ckpt_step*.pt"), key=lambda p: p.stat().st_mtime) + for p in cks[:-keep]: + try: p.unlink() + except: pass + +def safe_prompt(prompt: str, default: str) -> str: + try: + val = Prompt.ask(prompt, default=default) + return val + except (KeyboardInterrupt, EOFError): + console.print("\n[bold yellow]Input cancelled by user.[/] Using default.") + return default + +# ---- Loss CSV / Plot ---- + +def init_loss_csv(out_dir: str) -> str: + csv_path = os.path.join(out_dir, "loss_log.csv") + if not os.path.exists(csv_path): + with open(csv_path, "w", newline="", encoding="utf-8") as f: + w = csv.writer(f) + w.writerow(["step", "loss", "images_per_sec", "elapsed_sec"]) + return csv_path + +def append_loss(csv_path: str, step: int, loss_val: float, ips: float, elapsed: float): + try: + with open(csv_path, "a", newline="", encoding="utf-8") as f: + w = csv.writer(f) + w.writerow([step, f"{loss_val:.6f}", f"{ips:.3f}", f"{int(elapsed)}"]) + except Exception as e: + logging.warning(f"Could not write loss CSV: {e}") + +def try_plot_loss(csv_path: str, out_dir: str): + if not HAVE_MPL: + logging.info("matplotlib not installed; skipping loss plot.") + return + try: + steps, losses = [], [] + with open(csv_path, "r", encoding="utf-8") as f: + reader = csv.DictReader(f) + for row in reader: + steps.append(int(row["step"])) + losses.append(float(row["loss"])) + if steps and losses: + plt.figure() + plt.plot(steps, losses) + plt.xlabel("step") + plt.ylabel("loss (L1 + 0.1*aux)") + plt.title("Training Loss") + plt.tight_layout() + p = os.path.join(out_dir, "loss_curve.png") + plt.savefig(p, dpi=120) + plt.close() + logging.info(f"Saved loss plot: {p}") + except Exception as e: + logging.warning(f"Could not plot loss curve: {e}") + + +# ----------------- Datasets ----------------- + +class FolderPairs(Dataset): + def __init__(self, root: str, crop: int = 512, use_scribbles: bool = True): + self.root = root + self.crop = crop + self.use_scribbles = use_scribbles + exts = (".png", ".jpg", ".jpeg", ".webp", ".bmp") + self.paths = [str(Path(dp)/fn) for dp,_,files in os.walk(root) + for fn in files if fn.lower().endswith(exts)] + if not self.paths: + raise RuntimeError(f"No images found under: {root}") + + def __len__(self): return len(self.paths) + + def __getitem__(self, idx: int): + img = Image.open(self.paths[idx]).convert("RGB") + s = max(self.crop, min(img.size)) + img = img.resize((s, s), Image.BICUBIC) + import numpy as _np + x0 = _np.random.randint(0, s - self.crop + 1) + y0 = _np.random.randint(0, s - self.crop + 1) + img = img.crop((x0, y0, x0 + self.crop, y0 + self.crop)) + + rgb = np.asarray(img).astype(np.float32) / 255.0 + L = rgb_to_L(rgb) + + if self.use_scribbles: + hint_rgb, hint_mask = random_scribbles(self.crop, self.crop, n=np.random.randint(2, 5)) + hint_rgb = (hint_rgb.astype(np.float32) / 255.0 - 0.5) / 0.5 + hint_mask = (hint_mask.astype(np.float32) / 255.0)[..., None] + else: + hint_rgb = np.zeros((self.crop, self.crop, 3), dtype=np.float32) + hint_mask = np.zeros((self.crop, self.crop, 1), dtype=np.float32) + + inp = np.concatenate([L, hint_rgb * hint_mask, hint_mask], axis=2) + inp = torch.from_numpy(inp).permute(2, 0, 1) + tgt = torch.from_numpy(rgb).permute(2, 0, 1) * 2 - 1 + return inp, tgt + +class HFPairs(Dataset): + def __init__(self, dataset_id: str, split: str = "train", + crop: int = 512, streaming: bool = False, use_scribbles: bool = False): + if not HAVE_DATASETS: + raise RuntimeError("Please install: pip install datasets") + self.crop = crop + self.use_scribbles = use_scribbles + self.streaming = streaming + self.ds = load_dataset(dataset_id, split=split, streaming=streaming) + feats = getattr(self.ds, "features", None) + self.has_bw = bool(feats and "bw_image" in feats) + if streaming: + self._it = iter(self.ds) + + def __len__(self): return len(self.ds) if not self.streaming else 10**9 + + def _prep(self, bw_pil: Optional[Image.Image], color_pil: Image.Image): + color = color_pil.convert("RGB") + s = max(self.crop, min(color.size)) + color = color.resize((s, s), Image.BICUBIC) + import numpy as _np + x0 = _np.random.randint(0, s - self.crop + 1) + y0 = _np.random.randint(0, s - self.crop + 1) + color = color.crop((x0, y0, x0 + self.crop, y0 + self.crop)) + rgb = np.asarray(color).astype(np.float32) / 255.0 + + if bw_pil is None: + L = rgb_to_L(rgb) + else: + bw = bw_pil.convert("L").resize((s, s), Image.BICUBIC).crop((x0, y0, x0 + self.crop, y0 + self.crop)) + L = (np.asarray(bw, dtype=np.float32) / 255.0)[..., None] + + if self.use_scribbles: + hint_rgb, hint_mask = random_scribbles(self.crop, self.crop, n=np.random.randint(2, 5)) + hint_rgb = (hint_rgb.astype(np.float32) / 255.0 - 0.5) / 0.5 + hint_mask = (hint_mask.astype(np.float32) / 255.0)[..., None] + else: + hint_rgb = np.zeros((self.crop, self.crop, 3), dtype=np.float32) + hint_mask = np.zeros((self.crop, self.crop, 1), dtype=np.float32) + + inp = np.concatenate([L, hint_rgb * hint_mask, hint_mask], axis=2) + inp = torch.from_numpy(inp).permute(2, 0, 1) + tgt = torch.from_numpy(rgb).permute(2, 0, 1) * 2 - 1 + return inp, tgt + + def __getitem__(self, idx): + ex = next(self._it) if self.streaming else self.ds[int(idx)] + color = ex.get("color_image") or ex.get("image") + if color is None: + raise RuntimeError("Example missing 'color_image' (or 'image').") + bw = ex.get("bw_image") if self.has_bw else None + return self._prep(bw, color) + + +# ----------------- Training ----------------- + +def build_dataloader(source_type: str, crop: int, batch: int, workers: int, use_scribbles: bool): + if source_type == "folder": + folder = safe_prompt(f"[bold]Folder of COLOR images[/] [default {OUT_DIR}\\demo_images]", f"{OUT_DIR}\\demo_images") + ds = FolderPairs(folder, crop=crop, use_scribbles=use_scribbles) + desc = f"Folder: {folder} (N={len(ds)})" + else: + if not HAVE_DATASETS: + raise RuntimeError("Install datasets: pip install datasets") + ds_id = safe_prompt("HF dataset id", "MichaelP84/manga-colorization-dataset") + split = safe_prompt("Split", "train") + streaming = Confirm.ask("Enable streaming?", default=False) + ds = HFPairs(ds_id, split=split, crop=crop, streaming=streaming, use_scribbles=use_scribbles) + n = len(ds) if not streaming else "∞" + desc = f"HF: {ds_id} / {split} (N={n}, streaming={streaming})" + + # Windows-safe DataLoader defaults: + num_workers = 0 if IS_WINDOWS else max(0, workers) + pin_memory = (not IS_WINDOWS) + + dl_kwargs = dict( + batch_size=batch, + shuffle=True, + num_workers=num_workers, + pin_memory=pin_memory, + drop_last=True, + worker_init_fn=lambda _: set_seed() + ) + if (not IS_WINDOWS) and num_workers > 0: + dl_kwargs["prefetch_factor"] = 2 + dl_kwargs["persistent_workers"] = False + + dl = DataLoader(ds, **dl_kwargs) + return dl, desc + +def train_loop(cfg: dict): + ensure_dir(OUT_DIR) + ckpt_dir = os.path.join(OUT_DIR, "checkpoints"); ensure_dir(ckpt_dir) + sample_dir = os.path.join(OUT_DIR, "samples"); ensure_dir(sample_dir) + + # logging file + setup_logging(OUT_DIR) + with open(os.path.join(OUT_DIR, "run_config.json"), "w", encoding="utf-8") as f: + json.dump(cfg, f, indent=2) + + device = "cuda" if torch.cuda.is_available() else "cpu" + console.print(Panel.fit(device_info(device), border_style="cyan")) + if device != "cuda": + console.print("[yellow]CUDA not available. CPU training will be slow.[/]") + logging.info(mem_stats()) + + net = Colorizer().to(device) + sd = state_dict_load(GEN_ZIP) + missing, unexpected = net.generator.load_state_dict(sd, strict=False) + if missing: logging.info(f"Missing keys: {len(missing)}") + if unexpected: logging.info(f"Unexpected keys: {len(unexpected)}") + + # freeze encoder warmup + for p in net.generator.encoder.parameters(): p.requires_grad = False + opt = torch.optim.AdamW(filter(lambda p: p.requires_grad, net.parameters()), + lr=cfg["lr"], betas=(0.9,0.999), weight_decay=1e-4) + scaler = torch_amp.GradScaler(device="cuda" if device=="cuda" else "cpu", enabled=(device=="cuda")) + l1 = nn.L1Loss() + + dl, ds_desc = build_dataloader(cfg["source_type"], cfg["crop"], cfg["batch_size"], cfg["workers"], cfg["use_scribbles"]) + console.print(Panel.fit(f"[bold]Data:[/]\n{ds_desc}\n\n[dim]{mem_stats()}[/]")) + logging.info(ds_desc) + + # resume? + start_step = 0 + latest = latest_ckpt(ckpt_dir) + if latest and Confirm.ask(f"Resume from {latest.name}?", default=True): + ckpt = torch.load(latest, map_location="cpu") + net.load_state_dict(ckpt["model"]) + opt.load_state_dict(ckpt["opt"]) + scaler.load_state_dict(ckpt["scaler"]) + start_step = ckpt["step"] + logging.info(f"Resumed from step {start_step}") + + # init loss CSV + csv_path = init_loss_csv(OUT_DIR) + + # progress bar + progress = Progress( + TextColumn("[bold]Step[/] {task.completed}/{task.total}"), + BarColumn(), + TextColumn("loss {task.fields[loss]:.4f}"), + TextColumn("{task.fields[ips]}"), + TimeElapsedColumn(), + TextColumn("ETA"), + TimeRemainingColumn(), + console=console, + transient=False, + ) + task = progress.add_task("train", total=cfg["steps"], loss=0.0, ips="0.0 img/s") + + accum = max(1, cfg["grad_accum"]) + unfreeze_at = min(1000, cfg["steps"] // 5) + next_sample = cfg["sample_every"] + next_ckpt = cfg["save_every"] + grad_clip = cfg["grad_clip"] + + seen = 0 + start_time = time.time() + step = start_step + + try: + with progress: + while step < cfg["steps"]: + for inp, tgt in dl: + step += 1 + inp = inp.to(device, non_blocking=True) + tgt = tgt.to(device, non_blocking=True) + L = inp[:, :1]; hint4 = inp[:, 1:] + x = torch.cat([L, hint4], dim=1) + + # Modern autocast API + with torch_amp.autocast(device_type=("cuda" if device=="cuda" else "cpu"), enabled=(device=="cuda")): + pred, aux = net(x) + loss = l1(pred, tgt) + 0.1*l1(aux, tgt) + loss = loss / accum + + scaler.scale(loss).backward() + + if step % accum == 0: + if grad_clip: + scaler.unscale_(opt) + nn.utils.clip_grad_norm_(net.parameters(), grad_clip) + scaler.step(opt); scaler.update() + opt.zero_grad(set_to_none=True) + + # unfreeze + if step == unfreeze_at: + for p in net.generator.encoder.parameters(): p.requires_grad = True + for g in opt.param_groups: g["lr"] = cfg["lr"] * 0.25 + logging.info("Unfroze encoder; lowered LR") + logging.info(mem_stats()) + + # progress + CSV + seen += inp.size(0) + ips_float = seen / max(1e-6, time.time() - start_time) + ips = f"{ips_float:.1f} img/s" + progress.update(task, advance=1, loss=(loss.item()*accum), ips=ips) + append_loss(csv_path, step=step, loss_val=(loss.item()*accum), + ips=ips_float, elapsed=(time.time() - start_time)) + + # sample + if step >= next_sample or step == cfg["steps"]: + try: + with torch.no_grad(): + pv = (pred[0].clamp(-1,1).add(1).mul(0.5)).cpu().permute(1,2,0).numpy() + tv = (tgt[0].clamp(-1,1).add(1).mul(0.5)).cpu().permute(1,2,0).numpy() + lv = L[0,0].cpu().numpy(); lv = np.repeat(lv[...,None], 3, axis=2) + outp = os.path.join(OUT_DIR, "samples", f"sample_step{step}.jpg") + save_image_row(outp, [to_uint8(lv), to_uint8(pv), to_uint8(tv)]) + logging.info(f"Sample saved: {outp}") + except Exception as e: + logging.warning(f"Sample save failed: {e}") + next_sample += cfg["sample_every"] + + # ckpt + if step >= next_ckpt or step == cfg["steps"]: + ck = { + "step": step, + "model": net.state_dict(), + "opt": opt.state_dict(), + "scaler": scaler.state_dict(), + "config": cfg, + } + pth = os.path.join(OUT_DIR, "checkpoints", f"ckpt_step{step}.pt") + torch.save(ck, pth) + rotate_checkpoints(os.path.join(OUT_DIR, "checkpoints"), keep=cfg["keep_last"]) + logging.info(f"Checkpoint: {pth}") + next_ckpt += cfg["save_every"] + + if step >= cfg["steps"]: + break + + except KeyboardInterrupt: + console.print("\n[bold yellow]Training interrupted by user.[/]") + ck = { + "step": step, + "model": net.state_dict(), + "opt": opt.state_dict(), + "scaler": scaler.state_dict(), + "config": cfg, + } + pth = os.path.join(OUT_DIR, "checkpoints", f"ckpt_step{step}_INT.pt") + torch.save(ck, pth) + logging.info(f"Saved interrupt checkpoint: {pth}") + + finally: + # Export generator-only weights for your app + final_zip = os.path.join(OUT_DIR, f"generator_finetuned_step{step}.zip") + torch.save(net.generator.state_dict(), final_zip) + console.print(f"[bold green]Exported generator weights:[/] {final_zip}") + logging.info(f"Exported generator weights: {final_zip}") + + # Loss curve (optional) + try_plot_loss(csv_path, OUT_DIR) + + # ---- Compatibility export ---- + # Backup old app weight file, then drop-in replace it with the new one. + try: + ts = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") + if os.path.exists(GEN_ZIP): + bak = f"{GEN_ZIP}.{ts}.bak" + shutil.copy2(GEN_ZIP, bak) + logging.info(f"Backed up previous generator.zip -> {bak}") + shutil.copy2(final_zip, GEN_ZIP) + console.print(f"[bold green]Updated app weights:[/] {GEN_ZIP}") + logging.info(f"Updated app weights: {GEN_ZIP}") + except Exception as e: + logging.warning(f"Could not overwrite app weights: {e}") + console.print("[yellow]Note:[/] Could not overwrite app weights automatically. " + f"Copy {final_zip} to {GEN_ZIP} manually.") + + +# ----------------- Main (Rich prompts, graceful) ----------------- + +def main(): + console.print(Panel.fit("Manga Colorizer — Fine-tune", border_style="magenta")) + if not os.path.exists(GEN_ZIP): + console.print(f"[red]Model not found:[/] {GEN_ZIP}\nEdit GEN_ZIP at top of the script.") + return + + set_seed(42, deterministic=True) + + # Prompts (graceful + defaults; workers default to 0 on Windows) + src = safe_prompt("Data source — (1) Local folder (2) Hugging Face", "2") + source_type = "folder" if src.strip() == "1" else "hf" + + crop = int(safe_prompt("Crop size (divisible by 32)", "512")) + steps = int(safe_prompt("Total training steps", "4000")) + batch = int(safe_prompt("Batch size", "4")) + lr = float(safe_prompt("Learning rate", "2e-4")) + accum = int(safe_prompt("Grad accumulation (for bigger effective batch)", "1")) + default_workers = "0" if IS_WINDOWS else "2" + workers = int(safe_prompt("Dataloader workers", default_workers)) + save_every = int(safe_prompt("Save checkpoint every N steps", "500")) + sample_every = int(safe_prompt("Save sample image every N steps", "200")) + keep_last = int(safe_prompt("Keep last K checkpoints", "5")) + use_scribbles = Confirm.ask("Teach scribble hints too?", default=False) + gc = float(safe_prompt("Gradient clip (0 = off)", "0")) + grad_clip = None if gc <= 0 else gc + + cfg = { + "source_type": source_type, + "crop": crop, + "steps": steps, + "batch_size": batch, + "lr": lr, + "grad_accum": accum, + "workers": workers, + "save_every": save_every, + "sample_every": sample_every, + "keep_last": keep_last, + "use_scribbles": use_scribbles, + "grad_clip": grad_clip, + } + + # Show config table + tbl = Table(title="Run Config", show_header=False, box=None) + for k, v in cfg.items(): + tbl.add_row(f"[cyan]{k}[/]", f"{v}") + console.print(tbl) + + train_loop(cfg) + + +if __name__ == "__main__": + try: + main() + except KeyboardInterrupt: + console.print("\n[bold yellow]Exited by user during setup.[/]") + finally: + if PAUSE_ON_EXIT: + try: + input("\nDone. Press Enter to exit...") + except Exception: + pass diff --git a/upscalator.py b/upscalator.py new file mode 100644 index 0000000..7a16ee4 --- /dev/null +++ b/upscalator.py @@ -0,0 +1,51 @@ +import torch +import numpy as np + +from networks.RRDBNet import Upscaler as ESRUpscaler +from networks.aura_sr import Upscaler as GigaUpscaler, upscale_4x_overlapped, upscale_4x +from utils.utils import tile_process + +class MangaUpscaler: + def __init__(self, config): + if config.device == 'cuda' and not torch.cuda.is_available(): + print("[-] CUDA not available, using CPU.") + self.device = 'cpu' + else: + self.device = config.device + + self.tile_size = config.upscaler_tile_size + self.tile_pad = config.tile_pad + + if config.upscaler_type == 'GigaGAN': + self.model = GigaUpscaler().to(self.device) + else: + self.model = ESRUpscaler().to(self.device) + + model_or_chkpt = torch.load(config.upscaler_path, map_location=self.device, weights_only=False) + if config.upscaler_path.endswith(".pt"): + self.model.generator = model_or_chkpt + else: + self.model.generator.load_state_dict(model_or_chkpt, strict=True) + + self.model = self.model.eval() + + + def upscale(self, image, scale): + if image.shape[2] == 4: + image = image[:, :, :3] # Discard the alpha channel + + with torch.no_grad(): + if isinstance(self.model, GigaUpscaler): + result = upscale_4x(image, self.model) + else: + img_tensor = torch.from_numpy(image).to(self.device) + result = img_tensor.permute(2, 0, 1).unsqueeze(0) + if self.tile_size > 0: + result = tile_process(self.model, result.detach(), scale, self.tile_size, self.tile_pad) + else: + result = self.model(result.detach()) + result = result.data.squeeze().float().cpu().clamp_(0, 1).numpy() + result = np.transpose(result[[2, 1, 0], :, :], (1, 2, 0)) + result = (result * 255.0).round().astype(np.uint8) + result = result[:, :, ::-1] + return result diff --git a/utils/__init__.py b/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/utils/__pycache__/__init__.cpython-311.pyc b/utils/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000..b573746 Binary files /dev/null and b/utils/__pycache__/__init__.cpython-311.pyc differ diff --git a/utils/__pycache__/utils.cpython-311.pyc b/utils/__pycache__/utils.cpython-311.pyc new file mode 100644 index 0000000..53b4f15 Binary files /dev/null and b/utils/__pycache__/utils.cpython-311.pyc differ diff --git a/utils/utils.py b/utils/utils.py new file mode 100644 index 0000000..72c389c --- /dev/null +++ b/utils/utils.py @@ -0,0 +1,170 @@ +import base64 +import io +import math +import random +import re +import string + +import numpy as np +import cv2 +import torch +import PIL.ImageChops, PIL.ImageOps, PIL.Image + +try: + from matplotlib import pyplot as plt +except Exception: + plt = None + +def maybe_show(img, title="Preview"): + if plt is None: + return + plt.imshow(img) + plt.title(title) + plt.show() + +def resize_pad(img, size = 256): + + if len(img.shape) == 2: + img = np.expand_dims(img, 2) + + if img.shape[2] == 1: + img = np.repeat(img, 3, 2) + + if img.shape[2] == 4: + img = img[:, :, :3] + + pad = None + + if (img.shape[0] < img.shape[1]): + height = img.shape[0] + ratio = height / (size * 1.5) + width = int(np.ceil(img.shape[1] / ratio)) + img = cv2.resize(img, (width, int(size * 1.5)), interpolation = cv2.INTER_AREA) + + + new_width = width + (32 - width % 32) + + pad = (0, new_width - width) + + img = np.pad(img, ((0, 0), (0, pad[1]), (0, 0)), 'maximum') + else: + width = img.shape[1] + ratio = width / size + height = int(np.ceil(img.shape[0] / ratio)) + img = cv2.resize(img, (size, height), interpolation = cv2.INTER_AREA) + + new_height = height + (32 - height % 32) + + pad = (new_height - height, 0) + + img = np.pad(img, ((0, pad[0]), (0, 0), (0, 0)), 'maximum') + + if (img.dtype == 'float32'): + np.clip(img, 0, 1, out = img) + + return img[:, :, :1], pad + +def image_to_base64(img, format="WEBP"): + buffered = io.BytesIO() + img = PIL.Image.fromarray(img) + img.save(buffered, format=format) + buffered.seek(0) + img_byte = buffered.getvalue() + return f"data:image/{format.lower()};base64," + base64.b64encode(img_byte).decode('utf-8') + +def load_image_as_base64(filepath, format="WEBP"): + with open(filepath, "rb") as img_file: + img_byte = img_file.read() + return f"data:image/{format.lower()};base64," + base64.b64encode(img_byte).decode('utf-8') + +def save_image(image, filename, format="WEBP"): + image_pil = PIL.Image.fromarray(image) + image_pil.save(filename, format=format) + +def sanitize_string(input_string): + sanitized_string = re.sub(r'[^\w]', '_', input_string) + return sanitized_string + +def distance_from_grayscale(image): + try: + img_diff = PIL.ImageChops.difference(image, PIL.ImageOps.grayscale(image).convert('RGB')) + dist = np.array(img_diff.getdata()).mean() + return dist + except: + return 0 + +def generate_random_id(length=8): + characters = string.ascii_uppercase + string.digits + random_id = ''.join(random.choices(characters, k=length)) + return random_id + +def clear_torch_cache(): + torch.cuda.empty_cache() + +def tile_process(model, img, scale, tile_size, tile_pad): + if scale == 2: print('[-] ScaleFactor=2 is broken, please do not use it yet') + batch, channel, height, width = img.shape + output_height = height * scale + output_width = width * scale + output_shape = (batch, 3, output_height, output_width) + + # start with black image + output = img.new_zeros(output_shape) + tiles_x = math.ceil(width / tile_size) + tiles_y = math.ceil(height / tile_size) + + # loop over all tiles + for y in range(tiles_y): + for x in range(tiles_x): + # extract tile from input image + ofs_x = x * tile_size + ofs_y = y * tile_size + # input tile area on total image + input_start_x = ofs_x + input_end_x = min(ofs_x + tile_size, width) + input_start_y = ofs_y + input_end_y = min(ofs_y + tile_size, height) + + # input tile area on total image with padding + input_start_x_pad = max(input_start_x - tile_pad, 0) + input_end_x_pad = min(input_end_x + tile_pad, width) + input_start_y_pad = max(input_start_y - tile_pad, 0) + input_end_y_pad = min(input_end_y + tile_pad, height) + + # input tile dimensions + input_tile_width = input_end_x - input_start_x + input_tile_height = input_end_y - input_start_y + tile_idx = y * tiles_x + x + 1 + input_tile = img[:, :, input_start_y_pad:input_end_y_pad, input_start_x_pad:input_end_x_pad] + + # upscale tile + try: + with torch.no_grad(): + if model.name == 'colorizer': + output_tile,_ = model(input_tile) + # print(f'[+] Colorize Tile {tile_idx}/{tiles_x * tiles_y}') + + if model.name == 'upscaler': + output_tile = model(input_tile) + # print(f'[+] Upscale Tile {tile_idx}/{tiles_x * tiles_y}') + + except RuntimeError as error: + print('[!] Error: ', error) + + # output tile area on total image + output_start_x = input_start_x * scale + output_end_x = input_end_x * scale + output_start_y = input_start_y * scale + output_end_y = input_end_y * scale + + # output tile area without padding + output_start_x_tile = (input_start_x - input_start_x_pad) * scale + output_end_x_tile = output_start_x_tile + input_tile_width * scale + output_start_y_tile = (input_start_y - input_start_y_pad) * scale + output_end_y_tile = output_start_y_tile + input_tile_height * scale + + # put tile into output image + output[:, :, output_start_y:output_end_y, + output_start_x:output_end_x] = output_tile[:, :, output_start_y_tile:output_end_y_tile, + output_start_x_tile:output_end_x_tile] + return output \ No newline at end of file