diff --git a/Trainer/Trainer.bat b/Trainer/Trainer.bat deleted file mode 100644 index c184e6c..0000000 --- a/Trainer/Trainer.bat +++ /dev/null @@ -1,10 +0,0 @@ -@echo off -title Model Training -color 0b - -python Trainer.py - -echo. -echo Training Complete -echo. -pause > nul \ No newline at end of file diff --git a/Trainer/Trainer.py b/Trainer/Trainer.py deleted file mode 100644 index 38abf51..0000000 --- a/Trainer/Trainer.py +++ /dev/null @@ -1,414 +0,0 @@ -# Created by MSXYZ (AI-assisted) -# DLAA Model Trainer — v0.1.1 -# Perceptual Loss, Luma Loss, Smooth Loss, JPEG Augmentation - -# Steps: -# 1. Create a folder named "dataset" next to this script. -# 2. Add 50–200 high-quality images (landscapes, portraits, geometric patterns). -# 3. Run: python Trainer.py - -import torch -import torch.nn as nn -import torch.nn.functional as F -import torch.optim as optim -from torch.utils.data import DataLoader, Dataset -from torchvision import transforms -from torchvision.models import vgg16, VGG16_Weights -from PIL import Image, ImageFilter -import random -import io -import os -import glob - - -# Model -class DLAANet(nn.Module): - def __init__(self): - super().__init__() - - self.register_buffer("sobel_x", - torch.tensor([[-1,0,1],[-2,0,2],[-1,0,1]], dtype=torch.float32).view(1,1,3,3)) - self.register_buffer("sobel_y", - torch.tensor([[-1,-2,-1],[0,0,0],[1,2,1]], dtype=torch.float32).view(1,1,3,3)) - self.register_buffer("jitter_offsets", - torch.tensor([[0.25,0.25],[-0.25,-0.25],[-0.25,0.25],[0.25,-0.25]], dtype=torch.float32)) - - self.enc1 = nn.Sequential( - nn.Conv2d(3, 64, 3, padding=1, bias=False), - nn.GroupNorm(8, 64), - nn.LeakyReLU(0.2, inplace=True), - ) - self.enc2 = nn.Sequential( - nn.Conv2d(64, 64, 3, padding=1, bias=False), - nn.GroupNorm(8, 64), - nn.LeakyReLU(0.2, inplace=True), - ) - self.bottleneck = nn.Sequential( - nn.Conv2d(64, 64, 3, padding=2, dilation=2, bias=False), - nn.GroupNorm(8, 64), - nn.LeakyReLU(0.2, inplace=True), - - nn.Conv2d(64, 64, 3, padding=4, dilation=4, bias=False), - nn.GroupNorm(8, 64), - nn.LeakyReLU(0.2, inplace=True), - - nn.Conv2d(64, 64, 3, padding=2, dilation=2, bias=False), - nn.GroupNorm(8, 64), - nn.LeakyReLU(0.2, inplace=True), - ) - self.dec = nn.Sequential( - nn.Conv2d(128, 64, 3, padding=1, bias=False), - nn.GroupNorm(8, 64), - nn.LeakyReLU(0.2, inplace=True), - - nn.Conv2d(64, 32, 3, padding=1, bias=False), - nn.GroupNorm(4, 32), - nn.LeakyReLU(0.2, inplace=True), - ) - self.reconstructor = nn.Conv2d(32, 3, 3, padding=1, bias=False) - self._init_weights() - - def _init_weights(self): - for m in self.modules(): - if isinstance(m, nn.Conv2d): - nn.init.kaiming_normal_(m.weight, nonlinearity='leaky_relu') - elif isinstance(m, nn.GroupNorm): - nn.init.constant_(m.weight, 1) - nn.init.constant_(m.bias, 0) - nn.init.zeros_(self.reconstructor.weight) - - def forward(self, x): - e1 = self.enc1(x) - e2 = self.enc2(e1) - b = self.bottleneck(e2) - d = self.dec(torch.cat([b, e1], dim=1)) - r = self.reconstructor(d) - return x + r - - -# Perceptual Loss -class _PerceptualLoss(nn.Module): - def __init__(self, device): - super().__init__() - vgg = vgg16(weights=VGG16_Weights.DEFAULT).features[:23].to(device).eval() - for p in vgg.parameters(): - p.requires_grad = False - self.vgg = vgg - self.l1 = nn.L1Loss() - - # VGG normalization - mean = torch.tensor([0.485, 0.456, 0.406], device=device).view(1,3,1,1) - std = torch.tensor([0.229, 0.224, 0.225], device=device).view(1,3,1,1) - self.register_buffer("mean", mean) - self.register_buffer("std", std) - - def _normalize(self, x): - return (x - self.mean) / self.std - - def forward(self, pred, target): - p = self.vgg(self._normalize(pred)) - t = self.vgg(self._normalize(target)) - return self.l1(p, t) - - -# Combined Loss -class DLAALoss(nn.Module): - def __init__(self, device, - w_pixel=1.0, w_edge=0.5, w_freq=0.1, - w_luma=0.5, w_smooth=0.3, w_perceptual=0.2): - super().__init__() - self.l1 = nn.L1Loss() - self.perceptual = _PerceptualLoss(device) - self.w_pixel = w_pixel - self.w_edge = w_edge - self.w_freq = w_freq - self.w_luma = w_luma - self.w_smooth = w_smooth - self.w_perceptual = w_perceptual - - sx = torch.tensor([[-1,0,1],[-2,0,2],[-1,0,1]], dtype=torch.float32).view(1,1,3,3) - sy = torch.tensor([[-1,-2,-1],[0,0,0],[1,2,1]], dtype=torch.float32).view(1,1,3,3) - self._sx = sx - self._sy = sy - - def _luma(self, x): - return 0.2126*x[:,0:1] + 0.7152*x[:,1:2] + 0.0722*x[:,2:3] - - def _edge_map(self, x): - sx = self._sx.to(x.device) - sy = self._sy.to(x.device) - gray = self._luma(x) - ex = F.conv2d(gray, sx, padding=1) - ey = F.conv2d(gray, sy, padding=1) - return torch.sqrt(ex**2 + ey**2 + 1e-6) - - def _freq(self, x): - return torch.abs(torch.fft.fft2(self._luma(x))) - - def forward(self, pred, target): - # Pixel - pixel = self.l1(pred, target) - - # Edge - edge = self.l1(self._edge_map(pred), self._edge_map(target)) - - # Frequency - freq = self.l1(self._freq(pred), self._freq(target)) - - # Luma - luma = self.l1(self._luma(pred), self._luma(target)) - - # Smooth - edge_map = self._edge_map(target).detach() - smooth_mask = 1.0 - torch.sigmoid((edge_map - 0.1) * 20.0) - smooth = self.l1(pred * smooth_mask, target * smooth_mask) - - # Perceptual - perceptual = self.perceptual(pred, target) - - total = ( - self.w_pixel * pixel + - self.w_edge * edge + - self.w_freq * freq + - self.w_luma * luma + - self.w_smooth * smooth + - self.w_perceptual * perceptual - ) - return total, { - "pixel": pixel.item(), - "edge": edge.item(), - "luma": luma.item(), - "smooth": smooth.item(), - "perceptual": perceptual.item(), - } - - -# Dataset -class DLAADataset(Dataset): - SUPPORTED = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tiff"} - - def __init__(self, image_dir: str, size: int = 512): - self.image_paths = [ - p for p in glob.glob(os.path.join(image_dir, "*.*")) - if os.path.splitext(p)[1].lower() in self.SUPPORTED - ] - if not self.image_paths: - raise FileNotFoundError(f"No images found in '{image_dir}'.") - - self.size = size - self.to_tensor = transforms.ToTensor() - self.clean_tf = transforms.Compose([ - transforms.Resize((size, size), interpolation=transforms.InterpolationMode.LANCZOS), - ]) - - def __len__(self): - return len(self.image_paths) - - def _alias_nearest(self, img, factor): - s = int(self.size * factor) - return img.resize((s, s), Image.NEAREST).resize((self.size, self.size), Image.NEAREST) - - def _alias_bilinear_blur(self, img, factor): - s = int(self.size * factor) - img = img.resize((s, s), Image.BILINEAR).resize((self.size, self.size), Image.BILINEAR) - return img.filter(ImageFilter.GaussianBlur(radius=random.uniform(0.3, 0.9))) - - def _alias_checkerboard(self, img): - t = self.to_tensor(img) - t[:, ::2, 1::2] = 0 - t[:, 1::2, ::2] = 0 - t = F.avg_pool2d(t.unsqueeze(0), 3, stride=1, padding=1).squeeze(0) - return transforms.ToPILImage()(t.clamp(0, 1)) - - def _alias_mixed(self, img, factor): - s = int(self.size * factor) - img = img.resize((s, s), Image.NEAREST).resize((self.size, self.size), Image.BILINEAR) - return img.filter(ImageFilter.GaussianBlur(radius=random.uniform(0.5, 1.2))) - - def _alias_jpeg(self, img, quality): - buf = io.BytesIO() - img.save(buf, format="JPEG", quality=quality) - buf.seek(0) - return Image.open(buf).copy() - - def __getitem__(self, idx): - try: - img = Image.open(self.image_paths[idx]).convert("RGB") - except Exception: - img = Image.new("RGB", (self.size, self.size), (128, 128, 128)) - - clean = self.clean_tf(img) - mode = random.randint(0, 4) - factor = random.uniform(0.15, 0.45) - - if mode == 0: aliased = self._alias_nearest(clean, factor) - elif mode == 1: aliased = self._alias_bilinear_blur(clean, factor) - elif mode == 2: aliased = self._alias_checkerboard(clean) - elif mode == 3: aliased = self._alias_mixed(clean, factor) - else: aliased = self._alias_jpeg(clean, random.randint(40, 75)) - - return self.to_tensor(aliased), self.to_tensor(clean) - - -# Early Stop -class EarlyStopping: - def __init__(self, patience=12, min_delta=1e-5): - self.patience = patience - self.min_delta = min_delta - self.best = float("inf") - self.counter = 0 - - def step(self, loss): - if loss < self.best - self.min_delta: - self.best = loss - self.counter = 0 - else: - self.counter += 1 - return self.counter >= self.patience - - -# Training -def train( - dataset_dir = "dataset", - output_path = "DLAANet.pth", - image_size = 512, - batch_size = 4, - epochs = 100, - lr = 3e-4, - patience = 12, - w_pixel = 1.0, - w_edge = 0.8, - w_freq = 0.1, - w_luma = 0.8, - w_smooth = 0.3, - w_perceptual = 0.1, -): - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - print(f"[DLAA] Device : {device}") - - dataset = DLAADataset(dataset_dir, size=image_size) - dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, - num_workers=0, pin_memory=(device.type == "cuda")) - print(f"[DLAA] Dataset : {len(dataset)} images | {len(dataloader)} batches/epoch") - - model = DLAANet().to(device) - n_params = sum(p.numel() for p in model.parameters() if p.requires_grad) - print(f"[DLAA] Parameters : {n_params:,}") - - criterion = DLAALoss(device, - w_pixel=w_pixel, w_edge=w_edge, w_freq=w_freq, - w_luma=w_luma, w_smooth=w_smooth, w_perceptual=w_perceptual) - - optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4) - scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs, eta_min=5e-6) - stopper = EarlyStopping(patience=patience) - - best_loss = float("inf") - best_state = None - - print(f"\n{'='*62}") - print(f" Model Training v0.1.1 | {epochs} epochs | LR {lr:.0e}") - print(f" Loss: pixel={w_pixel} edge={w_edge} freq={w_freq} luma={w_luma} smooth={w_smooth} perceptual={w_perceptual}") - print(f"{'='*62}") - - for epoch in range(epochs): - model.train() - total_loss = 0.0 - sub_totals = {"pixel": 0., "edge": 0., "luma": 0., "smooth": 0., "perceptual": 0.} - - for inputs, targets in dataloader: - inputs = inputs.to(device) - targets = targets.to(device) - - optimizer.zero_grad() - outputs = torch.clamp(model(inputs), 0.0, 1.0) - loss, sub_dict = criterion(outputs, targets) - loss.backward() - torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0) - optimizer.step() - - total_loss += loss.item() - for k in sub_totals: - sub_totals[k] += sub_dict[k] - - scheduler.step() - - n = len(dataloader) - avg = total_loss / n - lr_now = scheduler.get_last_lr()[0] - marker = " ★" if avg < best_loss else "" - - # Detailed loss every 5 epochs - if (epoch + 1) % 5 == 0 or epoch == 0: - breakdown = " | " + " ".join( - f"{k}:{sub_totals[k]/n:.4f}" for k in sub_totals - ) - else: - breakdown = "" - - print(f"Epoch [{epoch+1:>3}/{epochs}] Loss: {avg:.5f} LR: {lr_now:.2e}{marker}{breakdown}") - - if avg < best_loss: - best_loss = avg - best_state = {k: v.clone() for k, v in model.state_dict().items()} - - # Writes the file to disk from each new record - torch.save(best_state, output_path) - print(f" [SYSTEM] Updated to the best model disk: {output_path} (Loss: {avg:.5f})") - - if stopper.step(avg): - print(f"\n[DLAA] Early stop at epoch {epoch+1} (no improvement for {patience} epochs)") - break - - save_state = best_state if best_state else model.state_dict() - torch.save(save_state, output_path) - print(f"\n[DLAA] ✓ Best model saved → {output_path} (loss: {best_loss:.5f})") - print(f"[DLAA] Replace DLAANet.pth in your ComfyUI node folder.") - - -if __name__ == "__main__": - try: - print("=" * 62) - print(" Model Training v0.1.1") - print("=" * 62) - - dataset_dir = "dataset" - if not os.path.isdir(dataset_dir): - print(f"[ERROR] '{dataset_dir}' folder not found.") - input("Press Enter to exit...") - exit(1) - - n_imgs = len([ - f for f in glob.glob(os.path.join(dataset_dir, "*.*")) - if os.path.splitext(f)[1].lower() in DLAADataset.SUPPORTED - ]) - - print(f"[INFO] Checking the dataset...") - if n_imgs < 20: - print(f"[WARN] Only {n_imgs} images found. 100+ recommended.") - else: - print(f"[INFO] {n_imgs} images found. Starting training...") - - train( - dataset_dir = "dataset", - output_path = "DLAANet.pth", - image_size = 512, - batch_size = 4, - epochs = 150, - lr = 3e-4, - patience = 20, - w_pixel = 0.5, - w_edge = 1.0, - w_freq = 0.5, - w_luma = 1.0, - w_smooth = 0.1, - w_perceptual = 0.5, - ) - - print("\nTraining finished successfully.") - input("Press Enter to exit...") - - except Exception as e: - print("\n[CRASH]") - print(e) - input("Press Enter to exit...") \ No newline at end of file diff --git a/Trainer/dataset/ReadMe.txt b/Trainer/dataset/ReadMe.txt deleted file mode 100644 index 69bc7ff..0000000 --- a/Trainer/dataset/ReadMe.txt +++ /dev/null @@ -1 +0,0 @@ -Put training images here \ No newline at end of file