Initialize

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)

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)

denoising/denoiser.py

@@ -0,0 +1,117 @@
+"""
+Denoise an image with the FFDNet denoising method
+
+Copyright (C) 2018, Matias Tassano <matias.tassano@parisdescartes.fr>
+
+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 <http://www.gnu.org/licenses/>.
+"""
+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)

denoising/functions.py

@@ -0,0 +1,101 @@
+"""
+Functions implementing custom NN layers
+
+Copyright (C) 2018, Matias Tassano <matias.tassano@parisdescartes.fr>
+
+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 <http://www.gnu.org/licenses/>.
+"""
+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

denoising/models.py

@@ -0,0 +1,100 @@
+"""
+Definition of the FFDNet model and its custom layers
+
+Copyright (C) 2018, Matias Tassano <matias.tassano@parisdescartes.fr>
+
+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 <http://www.gnu.org/licenses/>.
+"""
+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

denoising/utils.py

@@ -0,0 +1,66 @@
+"""
+Different utilities such as orthogonalization of weights, initialization of
+loggers, etc
+
+Copyright (C) 2018, Matias Tassano <matias.tassano@parisdescartes.fr>
+
+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 <http://www.gnu.org/licenses/>.
+"""
+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

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()

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

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("<Configure>",
+            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 \"<URL>\"\n"
+            # "  python -m mangadex_downloader \"<URL>\"\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 "<folder>_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()

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)

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))

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

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

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'

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

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

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

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