# DLAATexture Model Trainer
# Trains the 1x texture refinement model used by comfyui-solricks.
# Exports DLAATexture.pth and DLAATexture.safetensors.


import argparse
import io
import json
import math
import os
import random
import time
from pathlib import Path

import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
from safetensors.torch import save_file
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms
from torchvision.utils import save_image


try:
    from torchvision.models import vgg16, VGG16_Weights
    HAS_VGG = True
except Exception:
    HAS_VGG = False

try:
    from torch.amp import GradScaler, autocast
except Exception:
    from torch.cuda.amp import GradScaler, autocast


IMG_EXTS = {".png", ".jpg", ".jpeg", ".webp", ".bmp", ".tiff"}


class DLAANet(nn.Module):
    def __init__(self, base_channels=192):
        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.0000, -0.1667],
                [-0.2500, 0.1667],
                [0.2500, -0.3889],
                [-0.3750, -0.0556],
                [0.1250, 0.2778],
                [-0.1250, -0.2778],
                [0.3750, 0.0556],
                [-0.4375, 0.3889],
                [0.0625, -0.4630],
                [-0.3125, 0.1296],
                [0.1875, -0.2963],
                [-0.4375, -0.0370],
                [0.3125, 0.2407],
                [-0.0625, -0.2037],
                [0.4375, 0.0185],
                [-0.4688, 0.3148],
            ], dtype=torch.float32),
        )

        self.enc1 = nn.Sequential(
            nn.Conv2d(3, base_channels, 3, padding=1, bias=False),
            nn.GroupNorm(12, base_channels),
            nn.LeakyReLU(0.2, inplace=True),
        )
        self.enc2 = nn.Sequential(
            nn.Conv2d(base_channels, base_channels, 3, padding=1, bias=False),
            nn.GroupNorm(12, base_channels),
            nn.LeakyReLU(0.2, inplace=True),
        )
        self.bottleneck = nn.Sequential(
            nn.Conv2d(base_channels, base_channels, 3, padding=1, dilation=1, bias=False),
            nn.GroupNorm(12, base_channels),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(base_channels, base_channels, 3, padding=2, dilation=2, bias=False),
            nn.GroupNorm(12, base_channels),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(base_channels, base_channels, 3, padding=4, dilation=4, bias=False),
            nn.GroupNorm(12, base_channels),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(base_channels, base_channels, 3, padding=1, dilation=1, bias=False),
            nn.GroupNorm(12, base_channels),
            nn.LeakyReLU(0.2, inplace=True),
        )
        self.dec = nn.Sequential(
            nn.Conv2d(base_channels * 2, base_channels, 3, padding=1, bias=False),
            nn.GroupNorm(12, base_channels),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(base_channels, 64, 3, padding=1, bias=False),
            nn.GroupNorm(8, 64),
            nn.LeakyReLU(0.2, inplace=True),
        )
        self.reconstructor = nn.Conv2d(64, 3, 3, padding=1, bias=False)
        self._init_weights()

    def _init_weights(self):
        for layer in self.modules():
            if isinstance(layer, nn.Conv2d):
                nn.init.kaiming_normal_(layer.weight, nonlinearity="leaky_relu")
            elif isinstance(layer, nn.GroupNorm):
                nn.init.constant_(layer.weight, 1)
                nn.init.constant_(layer.bias, 0)
        nn.init.zeros_(self.reconstructor.weight)

    def forward(self, x):
        feat_low = self.enc1(x)
        feat_high = self.enc2(feat_low)
        context = self.bottleneck(feat_high)
        decoded = self.dec(torch.cat([context, feat_low], dim=1))
        residual = self.reconstructor(decoded)
        return torch.clamp(x + residual, 0.0, 1.0)


class TargetOnlyDataset(Dataset):
    def __init__(self, root, split="train", patch_size=192, augment=True, val_ratio=0.10, seed=42):
        self.root = Path(root)
        self.split = split
        self.patch_size = patch_size
        self.augment = augment and split == "train"

        if not self.root.exists():
            raise FileNotFoundError(f"Dataset folder not found: {self.root}")

        files = sorted([p for p in self.root.iterdir() if p.is_file() and p.suffix.lower() in IMG_EXTS])
        if not files:
            raise RuntimeError(f"No images found in: {self.root}")

        rng = random.Random(seed)
        rng.shuffle(files)

        val_count = max(1, int(len(files) * val_ratio)) if len(files) > 1 else 0
        val_files = files[:val_count]
        train_files = files[val_count:] if val_count > 0 else files

        self.files = val_files if split == "val" else train_files
        if not self.files:
            self.files = files

        self.to_tensor = transforms.ToTensor()

    def __len__(self):
        return len(self.files)

    def _load_rgb(self, path):
        return Image.open(path).convert("RGB")

    def _resize_if_needed(self, img):
        w, h = img.size
        ps = self.patch_size
        if w < ps or h < ps:
            scale = max(ps / max(w, 1), ps / max(h, 1))
            nw = max(ps, int(w * scale + 0.5))
            nh = max(ps, int(h * scale + 0.5))
            img = img.resize((nw, nh), Image.BICUBIC)
        return img

    def _random_crop(self, img):
        img = self._resize_if_needed(img)
        w, h = img.size
        ps = self.patch_size
        x = random.randint(0, w - ps)
        y = random.randint(0, h - ps)
        return img.crop((x, y, x + ps, y + ps))

    def _center_crop(self, img):
        img = self._resize_if_needed(img)
        w, h = img.size
        ps = self.patch_size
        x = max(0, (w - ps) // 2)
        y = max(0, (h - ps) // 2)
        return img.crop((x, y, x + ps, y + ps))

    def __getitem__(self, idx):
        path = self.files[idx]
        target = self._load_rgb(path)

        if self.split == "train":
            target = self._random_crop(target)
        else:
            target = self._center_crop(target)

        if self.augment:
            if random.random() < 0.5:
                target = target.transpose(Image.FLIP_LEFT_RIGHT)
            if random.random() < 0.5:
                target = target.transpose(Image.FLIP_TOP_BOTTOM)
            if random.random() < 0.25:
                target = target.rotate(random.choice([90, 180, 270]), expand=True)
                target = self._random_crop(target)

        target_t = self.to_tensor(target)
        input_a = degrade_input(target_t.unsqueeze(0), strong=True).squeeze(0)
        input_b = degrade_input(target_t.unsqueeze(0), strong=False).squeeze(0)

        return input_a, input_b, target_t, path.name


def _jpeg_roundtrip_tensor(x, quality):
    to_pil = transforms.ToPILImage()
    to_tensor = transforms.ToTensor()
    out = []
    for img in x:
        pil = to_pil(img.detach().cpu().clamp(0, 1))
        buf = io.BytesIO()
        pil.save(buf, format="JPEG", quality=quality)
        buf.seek(0)
        out.append(to_tensor(Image.open(buf).convert("RGB")))
    return torch.stack(out, dim=0).to(x.device)


def degrade_input(x, strong=True):
    b, c, h, w = x.shape
    device = x.device

    # Strong blur so the model learns to restore texture.
    if random.random() < 0.8:
        sigma = random.uniform(0.5, 2.0) if strong else random.uniform(0.2, 1.0)
        x = transforms.functional.gaussian_blur(x, kernel_size=(5,5), sigma=sigma)

    # Resolution loss through downscale and upscale.
    scale = random.uniform(0.4, 0.6) if strong else random.uniform(0.7, 0.9)
    nh, nw = max(16, int(h * scale)), max(16, int(w * scale))
    
    # Pick a random downscale mode for varied texture loss.
    down_mode = random.choice(["bilinear", "bicubic", "area"])
    y = F.interpolate(x, size=(nh, nw), mode=down_mode, align_corners=False if down_mode != "area" else None)
    
    # Upscale back to the original size, leaving a slightly blurred image.
    y = F.interpolate(y, size=(h, w), mode="bilinear", align_corners=False)

    # Light noise so the model sees texture under mild degradation.
    if random.random() < 0.3:
        noise = torch.randn_like(y) * random.uniform(0.001, 0.005)
        y = torch.clamp(y + noise, 0.0, 1.0)

    return torch.clamp(y, 0.0, 1.0)


class PerceptualLoss(nn.Module):
    def __init__(self, device):
        super().__init__()
        if not HAS_VGG:
            self.vgg = None
            return
        try:
            vgg = vgg16(weights=VGG16_Weights.DEFAULT).features[:8].to(device).eval()
        except Exception:
            self.vgg = None
            return
        for p in vgg.parameters():
            p.requires_grad = False
        self.vgg = vgg
        self.register_buffer("mean", torch.tensor([0.485, 0.456, 0.406], device=device).view(1, 3, 1, 1))
        self.register_buffer("std", torch.tensor([0.229, 0.224, 0.225], device=device).view(1, 3, 1, 1))

    def forward(self, pred, target):
        if self.vgg is None:
            return pred.new_tensor(0.0)
        pred_n = (pred - self.mean) / self.std
        tgt_n = (target - self.mean) / self.std
        return F.l1_loss(self.vgg(pred_n), self.vgg(tgt_n))


class LossPack(nn.Module):
    def __init__(self, device, use_perceptual=True):
        super().__init__()
        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)
        lap = torch.tensor([[0, 1, 0], [1, -4, 1], [0, 1, 0]], dtype=torch.float32).view(1, 1, 3, 3)
        self.register_buffer("sx", sx)
        self.register_buffer("sy", sy)
        self.register_buffer("lap", lap)
        self.perceptual = PerceptualLoss(device) if use_perceptual else None

    def luma(self, x):
        return 0.299 * x[:, 0:1] + 0.587 * x[:, 1:2] + 0.114 * x[:, 2:3]

    def edge_map(self, x):
        y = self.luma(x)
        ex = F.conv2d(y, self.sx, padding=1)
        ey = F.conv2d(y, self.sy, padding=1)
        return torch.sqrt(ex * ex + ey * ey + 1e-6)

    def lap_map(self, x):
        return F.conv2d(self.luma(x), self.lap, padding=1)

    def charbonnier(self, pred, target, eps=1e-3):
        diff = pred - target
        return torch.mean(torch.sqrt(diff * diff + eps * eps))

    def frequency_loss(self, pred, target):
        pred_fft = torch.fft.rfft2(pred, norm="ortho")
        tgt_fft = torch.fft.rfft2(target, norm="ortho")
        return F.l1_loss(torch.log1p(torch.abs(pred_fft)), torch.log1p(torch.abs(tgt_fft)))

    def smooth_loss(self, x):
        dx = torch.mean(torch.abs(x[:, :, :, :-1] - x[:, :, :, 1:]))
        dy = torch.mean(torch.abs(x[:, :, :-1, :] - x[:, :, 1:, :]))
        return dx + dy

    def forward(self, pred_a, pred_b, target, args):
        l1 = F.l1_loss(pred_a, target)
        edge = F.l1_loss(self.edge_map(pred_a), self.edge_map(target))
        lap = F.l1_loss(self.lap_map(pred_a), self.lap_map(target))
        freq = self.frequency_loss(pred_a, target)
        charb = self.charbonnier(pred_a, target)
        smooth = self.smooth_loss(pred_a)
        temporal = F.l1_loss(pred_a, pred_b)

        if self.perceptual is not None:
            perceptual = self.perceptual(pred_a, target)
        else:
            perceptual = pred_a.new_tensor(0.0)

        total = (
            args.w_l1 * l1 +
            args.w_edge * edge +
            args.w_lap * lap +
            args.w_freq * freq +
            args.w_charb * charb +
            args.w_smooth * smooth +
            args.w_temporal * temporal +
            args.w_perceptual * perceptual
        )

        parts = {
            "total": float(total.detach().cpu()),
            "l1": float(l1.detach().cpu()),
            "edge": float(edge.detach().cpu()),
            "lap": float(lap.detach().cpu()),
            "freq": float(freq.detach().cpu()),
            "charb": float(charb.detach().cpu()),
            "smooth": float(smooth.detach().cpu()),
            "temporal": float(temporal.detach().cpu()),
            "perceptual": float(perceptual.detach().cpu()),
        }
        return total, parts


def make_fp16_state_dict(state_dict):
    packed = {}
    for k, v in state_dict.items():
        if torch.is_floating_point(v):
            packed[k] = v.detach().cpu().half()
        else:
            packed[k] = v.detach().cpu()
    return packed


@torch.no_grad()
def validate(model, loader, loss_fn, device, args, max_batches=None):
    model.eval()
    total = 0.0
    count = 0
    for bi, (inp_a, inp_b, tgt, _) in enumerate(loader):
        inp_a = inp_a.to(device, non_blocking=True)
        inp_b = inp_b.to(device, non_blocking=True)
        tgt = tgt.to(device, non_blocking=True)
        pred_a = model(inp_a)
        pred_b = model(inp_b)
        loss, _ = loss_fn(pred_a, pred_b, tgt, args)
        total += float(loss.detach().cpu())
        count += 1
        if max_batches is not None and bi + 1 >= max_batches:
            break
    model.train()
    return total / max(count, 1)


@torch.no_grad()
def save_preview(model, loader, device, out_dir, step, max_items=4):
    model.eval()
    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    inp_a, inp_b, tgt, _ = next(iter(loader))
    inp_a = inp_a.to(device)[:max_items]
    inp_b = inp_b.to(device)[:max_items]
    tgt = tgt.to(device)[:max_items]
    pred = model(inp_a)
    grid = torch.cat([inp_a, inp_b, pred, tgt], dim=0)
    save_image(grid, out_dir / f"preview_{step:06d}.png", nrow=max_items)
    model.train()


def model_size_mb(path):
    return os.path.getsize(path) / (1024 * 1024)


def save_checkpoint(model, optimizer, scheduler, epoch, step, val_loss, out_dir, name):
    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    payload = {
        "epoch": epoch,
        "step": step,
        "val_loss": val_loss,
        "model": model.state_dict(),
        "optimizer": optimizer.state_dict(),
        "scheduler": scheduler.state_dict() if scheduler is not None else None,
    }
    pth_path = out_dir / f"{name}.pth"
    torch.save(payload, pth_path)
    save_file(model.state_dict(), str(out_dir / f"{name}.safetensors"))
    return model_size_mb(pth_path)


def save_best_snapshot(model, optimizer, scheduler, epoch, step, val_loss, out_dir, export_fp16=True):
    archive_dir = Path(out_dir) / "archive"
    archive_dir.mkdir(parents=True, exist_ok=True)
    clean_val = f"{val_loss:.6f}".replace(".", "_")
    name = f"best_epoch_{epoch:03d}_val_{clean_val}"
    payload = {
        "epoch": epoch,
        "step": step,
        "val_loss": val_loss,
        "model": model.state_dict(),
        "optimizer": optimizer.state_dict(),
        "scheduler": scheduler.state_dict() if scheduler is not None else None,
    }
    torch.save(payload, archive_dir / f"{name}.pth")
    state = make_fp16_state_dict(model.state_dict()) if export_fp16 else model.state_dict()
    save_file(state, str(archive_dir / f"{name}.safetensors"))
    return name


def save_inference_weights(model, out_dir, name="DLAATexture", export_fp16=True):
    out_dir = Path(out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    state = make_fp16_state_dict(model.state_dict()) if export_fp16 else model.state_dict()

    pth_path = out_dir / f"{name}.pth"
    safe_path = out_dir / f"{name}.safetensors"

    torch.save(state, pth_path)
    save_file(state, str(safe_path))

    return model_size_mb(pth_path), model_size_mb(safe_path)


def train(args):
    random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)

    device = torch.device("cuda" if torch.cuda.is_available() and not args.cpu else "cpu")
    out_dir = Path(args.out_dir)
    ckpt_dir = out_dir / "checkpoints"
    preview_dir = out_dir / "previews"
    out_dir.mkdir(parents=True, exist_ok=True)
    ckpt_dir.mkdir(parents=True, exist_ok=True)
    preview_dir.mkdir(parents=True, exist_ok=True)

    train_set = TargetOnlyDataset(args.dataset, "train", args.patch_size, augment=True, val_ratio=args.val_ratio, seed=args.seed)
    val_set = TargetOnlyDataset(args.dataset, "val", args.patch_size, augment=False, val_ratio=args.val_ratio, seed=args.seed)

    train_loader = DataLoader(
        train_set,
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=args.workers,
        pin_memory=(device.type == "cuda"),
        drop_last=True,
    )
    val_loader = DataLoader(
        val_set,
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=max(0, args.workers // 2),
        pin_memory=(device.type == "cuda"),
        drop_last=False,
    )

    model = DLAANet(base_channels=args.base_channels).to(device)
    loss_fn = LossPack(device, use_perceptual=not args.no_perceptual).to(device)

    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs, eta_min=args.lr * 0.05)
    scaler = GradScaler(enabled=args.amp and device.type == "cuda")

    n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    best_val = math.inf
    step = 0
    start_epoch = 1
    log_data = []
    log_path = out_dir / "training_log.json"

    if log_path.exists():
        try:
            log_data = json.loads(log_path.read_text(encoding="utf-8"))
        except Exception:
            log_data = []

    if args.resume is not None:
        resume_path = Path(args.resume)

        if not resume_path.exists():
            raise FileNotFoundError(f"Resume checkpoint not found: {resume_path}")

        checkpoint = torch.load(resume_path, map_location=device)

        model.load_state_dict(checkpoint["model"], strict=False)

        if "optimizer" in checkpoint and checkpoint["optimizer"] is not None:
            optimizer.load_state_dict(checkpoint["optimizer"])

        if "scheduler" in checkpoint and checkpoint["scheduler"] is not None:
            scheduler.load_state_dict(checkpoint["scheduler"])

        start_epoch = int(checkpoint.get("epoch", 0)) + 1
        step = int(checkpoint.get("step", 0))
        best_val = float(checkpoint.get("val_loss", math.inf))

        best_path = ckpt_dir / "best.pth"
        if best_path.exists():
            try:
                best_checkpoint = torch.load(best_path, map_location="cpu")
                best_val = float(best_checkpoint.get("val_loss", best_val))
            except Exception:
                pass

        print("==============================================================")
        print(f"Resuming from : {resume_path}")
        print(f"Start epoch   : {start_epoch}")
        print(f"Step          : {step}")
        print(f"Best val      : {best_val:.6f}")
        print("==============================================================")

    print("==============================================================")
    print("DLAATexture Trainer")
    print("==============================================================")
    print(f"Device      : {device}")
    print(f"Dataset     : {len(train_set)} train | {len(val_set)} val")
    print(f"Patch       : {args.patch_size}")
    print(f"Batch       : {args.batch_size}")
    print(f"Epochs      : {args.epochs}")
    print(f"Parameters  : {n_params / 1e6:.2f}M")
    print(f"Export      : {'FP16' if args.export_fp16 else 'FP32'}")
    print("==============================================================")

    for epoch in range(start_epoch, args.epochs + 1):
        model.train()
        epoch_loss = 0.0
        t0 = time.time()

        for inp_a, inp_b, tgt, _ in train_loader:
            inp_a = inp_a.to(device, non_blocking=True)
            inp_b = inp_b.to(device, non_blocking=True)
            tgt = tgt.to(device, non_blocking=True)

            optimizer.zero_grad(set_to_none=True)

            with autocast(device_type="cuda", enabled=args.amp and device.type == "cuda"):
                pred_a = model(inp_a)
                pred_b = model(inp_b)
                loss, parts = loss_fn(pred_a, pred_b, tgt, args)

            scaler.scale(loss).backward()
            if args.grad_clip > 0:
                scaler.unscale_(optimizer)
                torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
            scaler.step(optimizer)
            scaler.update()

            epoch_loss += float(loss.detach().cpu())
            step += 1

            if step % args.log_every == 0:
                lr = optimizer.param_groups[0]["lr"]
                print(
                    f"[E{epoch:03d} S{step:06d}] loss={parts['total']:.6f} "
                    f"l1={parts['l1']:.5f} edge={parts['edge']:.5f} "
                    f"freq={parts['freq']:.5f} temp={parts['temporal']:.5f} lr={lr:.2e}"
                )

        scheduler.step()

        train_loss = epoch_loss / max(len(train_loader), 1)
        val_loss = validate(model, val_loader, loss_fn, device, args, max_batches=args.val_batches)

        if epoch % args.preview_every == 0 or epoch == 1:
            save_preview(model, val_loader, device, preview_dir, step)

        last_size = save_checkpoint(model, optimizer, scheduler, epoch, step, val_loss, ckpt_dir, "last")

        is_best = val_loss < best_val
        if is_best:
            best_val = val_loss
            save_checkpoint(model, optimizer, scheduler, epoch, step, val_loss, ckpt_dir, "best")
            snapshot_name = save_best_snapshot(model, optimizer, scheduler, epoch, step, val_loss, ckpt_dir, export_fp16=args.export_fp16)
            
            pth_size, safe_size = save_inference_weights(
                model,
                ckpt_dir,
                "DLAATexture",
                export_fp16=args.export_fp16
            )
            
            if pth_size > args.max_model_mb:
                print(f"[WARN] DLAATexture.pth is {pth_size:.2f} MB, above limit {args.max_model_mb:.2f} MB")

            if safe_size > args.max_model_mb:
                print(f"[WARN] DLAATexture.safetensors is {safe_size:.2f} MB, above limit {args.max_model_mb:.2f} MB")

            print(
                f"[BEST] val={best_val:.6f} | "
                f"DLAATexture.pth={pth_size:.2f} MB | "
                f"DLAATexture.safetensors={safe_size:.2f} MB | "
                f"archived={snapshot_name}"
            )

        elapsed = time.time() - t0
        row = {
            "epoch": epoch,
            "step": step,
            "train_loss": train_loss,
            "val_loss": val_loss,
            "best_val": best_val,
            "lr": optimizer.param_groups[0]["lr"],
            "seconds": elapsed,
            "last_checkpoint_mb": last_size,
        }
        log_data.append(row)
        log_path.write_text(json.dumps(log_data, indent=2), encoding="utf-8")

        best_mark = " ★" if is_best else ""
        print(
            f"[E{epoch:03d}] train={train_loss:.6f} "
            f"val={val_loss:.6f} best={best_val:.6f} "
            f"time={elapsed:.1f}s{best_mark}"
        )

    print("Training complete.")
    print(f"Best validation loss: {best_val:.6f}")
    print(f"Output: {out_dir}")


def parse_args():
    parser = argparse.ArgumentParser(description="DLAATexture Trainer")
    parser.add_argument("--dataset", type=str, default="dataset")
    parser.add_argument("--batch-size", type=int, default=8)
    parser.add_argument("--patch-size", type=int, default=192)
    parser.add_argument("--lr", type=float, default=1e-4)
    parser.add_argument("--weight-decay", type=float, default=1e-4)
    parser.add_argument("--base-channels", type=int, default=192)
    parser.add_argument("--workers", type=int, default=4)
    parser.add_argument("--amp", action="store_true", default=True)
    parser.add_argument("--cpu", action="store_true")
    parser.add_argument("--grad-clip", type=float, default=1.0)
    parser.add_argument("--log-every", type=int, default=50)
    parser.add_argument("--preview-every", type=int, default=1)
    parser.add_argument("--val-batches", type=int, default=None)
    parser.add_argument("--max-model-mb", type=float, default=9.0)
    parser.add_argument("--val-ratio", type=float, default=0.10)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--export-fp16", action="store_true", default=True)
    parser.add_argument("--no-perceptual", action="store_true")
    parser.add_argument("--out-dir", type=str, default="runs")
    parser.add_argument("--epochs", type=int, default=120)
    parser.add_argument("--resume", type=str, default=None)

    parser.add_argument("--w-l1", type=float, default=1.00)
    parser.add_argument("--w-edge", type=float, default=0.5)
    parser.add_argument("--w-lap", type=float, default=0.30)
    parser.add_argument("--w-freq", type=float, default=0.5)
    parser.add_argument("--w-charb", type=float, default=0.12)
    parser.add_argument("--w-smooth", type=float, default=0.01)
    parser.add_argument("--w-temporal", type=float, default=0.2)
    parser.add_argument("--w-perceptual", type=float, default=0.4)
    return parser.parse_args()


if __name__ == "__main__":
    train(parse_args())