#!/usr/bin/env python3 """ scripts/train_demorpher.py Training script for the MorphGuard demorpher models: 1. Transformer-based demorpher (default) 2. GAN-based demorpher (with the --gan flag) Usage for transformer-based demorpher: python scripts/train_demorpher.py \ --data-dir data \ --epochs 20 \ --batch-size 16 \ --lr 1e-4 \ --num-pairs 10000 \ --save-path models/demorpher.pth \ [--gpus 1] [--num-workers 4] Usage for GAN-based demorpher: python scripts/train_demorpher.py \ --gan \ --data-dir data \ --epochs 30 \ --batch-size 4 \ --lr 1e-4 \ --model-type toonify \ --save-path models/demorpher_gan.pth """ import os import sys # ensure project root is on PYTHONPATH so imports like morphguard_api work sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) import argparse import random from PIL import Image import numpy as np import torch import torch.nn as nn from torch.utils.data import Dataset, DataLoader from torchvision import transforms # Import the demorpher module definition from morphguard_api import DemorphingModule class DemorphDataset(Dataset): """Dataset of randomly synthesized morph pairs from real face images.""" def __init__(self, data_dir, split='train', transform_input=None, transform_target=None, num_pairs=10000): self.real_dir = os.path.join(data_dir, split, 'real') self.files = [f for f in os.listdir(self.real_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))] if len(self.files) < 2: raise RuntimeError(f"Not enough real images in {self.real_dir} to generate morphs") self.transform_input = transform_input self.transform_target = transform_target self.num_pairs = num_pairs def __len__(self): return self.num_pairs def __getitem__(self, idx): # Sample two distinct real images a, b = random.sample(self.files, 2) path_a = os.path.join(self.real_dir, a) path_b = os.path.join(self.real_dir, b) img_a = Image.open(path_a).convert('RGB').resize((224, 224)) img_b = Image.open(path_b).convert('RGB').resize((224, 224)) # Create morph by averaging pixel values arr_a = np.array(img_a).astype(np.float32) arr_b = np.array(img_b).astype(np.float32) morph_arr = ((arr_a + arr_b) / 2.0).astype(np.uint8) img_morph = Image.fromarray(morph_arr) # Apply transforms if self.transform_input: morph_tensor = self.transform_input(img_morph) ref_tensor = self.transform_input(img_a) else: morph_tensor = transforms.ToTensor()(img_morph) ref_tensor = transforms.ToTensor()(img_a) if self.transform_target: target_tensor = self.transform_target(img_b) else: target_tensor = transforms.ToTensor()(img_b) return morph_tensor, ref_tensor, target_tensor def train_demorpher(args): # Determine device use_cuda = args.gpus > 0 and torch.cuda.is_available() device = torch.device('cuda' if use_cuda else 'cpu') # Define transforms for inputs and targets transform_input = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ]) transform_target = transforms.Compose([ transforms.Resize((224, 224)), transforms.ToTensor() ]) # Prepare datasets and loaders train_ds = DemorphDataset( args.data_dir, split='train', transform_input=transform_input, transform_target=transform_target, num_pairs=args.num_pairs ) val_pairs = max(1, int(args.num_pairs * args.val_ratio)) val_ds = DemorphDataset( args.data_dir, split='val', transform_input=transform_input, transform_target=transform_target, num_pairs=val_pairs ) train_loader = DataLoader( train_ds, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, pin_memory=use_cuda ) val_loader = DataLoader( val_ds, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=use_cuda ) # Initialize model, loss, optimizer model = DemorphingModule().to(device) criterion = nn.MSELoss() optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) best_val_loss = float('inf') # Setup statistics logging import json # noqa: E402 stats = {'epochs': [], 'train_loss': [], 'val_loss': []} # Training loop for epoch in range(1, args.epochs + 1): model.train() total_train = 0.0 for morph, ref, target in train_loader: morph = morph.to(device) ref = ref.to(device) target = target.to(device) optimizer.zero_grad() output = model(morph, ref) loss = criterion(output, target) loss.backward() optimizer.step() total_train += loss.item() * morph.size(0) avg_train = total_train / len(train_loader.dataset) model.eval() total_val = 0.0 with torch.no_grad(): for morph, ref, target in val_loader: morph = morph.to(device) ref = ref.to(device) target = target.to(device) output = model(morph, ref) loss = criterion(output, target) total_val += loss.item() * morph.size(0) avg_val = total_val / len(val_loader.dataset) print(f"Epoch {epoch}/{args.epochs} | train_loss={avg_train:.4f} | val_loss={avg_val:.4f}") # Update and write stats stats['epochs'].append(epoch) stats['train_loss'].append(avg_train) stats['val_loss'].append(avg_val) try: with open(args.stats_file, 'w') as sf: json.dump(stats, sf, indent=2) except Exception as e: print(f"Warning: could not write stats file: {e}") # Save best model if avg_val < best_val_loss: best_val_loss = avg_val os.makedirs(os.path.dirname(args.save_path), exist_ok=True) torch.save(model.state_dict(), args.save_path) print(f"Saved best demorpher model to {args.save_path} (val_loss={avg_val:.4f})") def parse_args(): parser = argparse.ArgumentParser(description='Train MorphGuard Demorpher') # Common parameters parser.add_argument('--data-dir', type=str, default='data', help='Root data directory') parser.add_argument('--epochs', type=int, default=20, help='Number of training epochs') parser.add_argument('--batch-size', type=int, default=16, help='Batch size') parser.add_argument('--lr', type=float, default=1e-4, help='Learning rate') parser.add_argument('--gpus', type=int, default=1, help='GPUs to use (0 for CPU)') parser.add_argument('--num-workers', type=int, default=4, help='DataLoader workers') parser.add_argument('--save-path', type=str, default='models/demorpher.pth', help='Path to save trained model') parser.add_argument('--stats-file', type=str, default='training_demorpher_stats.json', help='JSON file to write metrics') # Parameters for transformer-based model parser.add_argument('--num-pairs', type=int, default=10000, help='Number of morph pairs per epoch (transformer model)') parser.add_argument('--val-ratio', type=float, default=0.1, help='Fraction of pairs for validation (transformer model)') # Parameters for GAN-based model parser.add_argument('--gan', action='store_true', help='Train GAN-based demorpher instead of transformer-based') parser.add_argument('--model-type', type=str, default='toonify', choices=['toonify', 'frontalization', 'superresolution', 'sketch_to_face', 'inversion'], help='GAN model type to train (when using --gan)') parser.add_argument('--checkpoint-path', type=str, default=None, help='Path to specific checkpoint for GAN fine-tuning (when using --gan)') parser.add_argument('--output-dir', type=str, default='models/demorpher', help='Directory to save GAN checkpoints and logs (when using --gan)') return parser.parse_args() if __name__ == '__main__': args = parse_args() # If using GAN-based demorpher, use a different training function if args.gan: try: # Import GAN training modules sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) from plugins.gan_demorpher import GanDemorpher, AVAILABLE_MODEL_TYPES # Before importing pSp Coach, check if pSp is available try: # Import pSp modules psp_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'pSp')) sys.path.insert(0, psp_path) from pSp.training.coach import Coach from pSp.options.train_options import TrainOptions # Import the GAN training functionality from a separate module import importlib.util spec = importlib.util.spec_from_file_location( "train_gan_demorpher", os.path.abspath(os.path.join(os.path.dirname(__file__), 'train_gan_demorpher.py')) ) gan_module = importlib.util.module_from_spec(spec) spec.loader.exec_module(gan_module) # Use the GAN training function print(f"Training GAN-based demorpher with model type: {args.model_type}") gan_module.train_gan_demorpher(args) except ImportError as e: print(f"Error: pixel2style2pixel (pSp) not found or not properly configured: {e}") print("Please ensure the pSp subdirectory is correctly set up for GAN-based training.") sys.exit(1) except ImportError as e: print(f"Error: GAN demorpher module not available: {e}") print("Falling back to transformer-based demorpher training.") train_demorpher(args) else: # Use the regular transformer-based demorpher training train_demorpher(args)