mirror of
https://github.com/Nighthawk42/manga-colorizer.git
synced 2026-08-30 07:22:28 +00:00
594 lines
23 KiB
Python
594 lines
23 KiB
Python
# 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
|