scripts/train_demorpher.py

#!/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)