train.py

import os
import shutil
import argparse
import yaml
from tqdm import tqdm
import torch
import utils.misc as utils
import numpy as np
import wandb
from PIL import Image

from loss import TextureLoss
from models import NCA, NoiseNCA, PENCA

parser = argparse.ArgumentParser()
parser.add_argument('--config', type=str, default='configs/Noise-NCA.yml', help="Path to the config file")
parser.add_argument('--data_dir', type=str, default='data/images/texture/', help="Texture images directory")

os.environ["WANDB_SILENT"] = "true"


def get_nca_model(config, texture_name):
    model_type = config['model']['type']
    if model_type == 'NCA':
        return NCA(**config['model']['attr'])
    elif model_type == 'NoiseNCA':
        noise_levels = config['model']['noise_levels']
        if texture_name in noise_levels:
            noise_level = noise_levels[texture_name]
        else:
            noise_level = noise_levels['default']
        return NoiseNCA(noise_level=noise_level, **config['model']['attr'])
    elif model_type == 'PENCA':
        return PENCA(**config['model']['attr'])
    else:
        raise ValueError(f"Unknown model type: {model_type}")


def main(config):
    wandb_enabled = 'wandb' in config
    if wandb_enabled:
        wandb.login(key=config['wandb']['key'], relogin=True)

    # Auto-detect device: use CUDA if available, otherwise fall back to CPU
    requested_device = config.get('device', 'cuda:0')
    if 'cuda' in requested_device.lower() and not torch.cuda.is_available():
        print(f"Warning: CUDA requested but not available. Falling back to CPU.")
        device = torch.device('cpu')
    else:
        device = torch.device(requested_device)
    
    config['loss']['attr']['device'] = device
    config['model']['attr']['device'] = device
    loss_fn = TextureLoss(**config['loss']['attr']).to(device)

    data_path = config['data_dir']
    
    # Check if data_path is a file or directory
    if os.path.isfile(data_path):
        # Single image file: train only on this file
        image_paths = [data_path]
    elif os.path.isdir(data_path):
        # Directory: train on all files in directory (original behavior)
        image_paths = [f"{data_path}/{f}" for f in os.listdir(data_path)]
    else:
        raise ValueError(f"data_dir must be either a file or directory: {data_path}")

    for idx, url in enumerate(image_paths):
        if "ipynb" in url:
            continue

        style_img = utils.imread(url, max_size=128)
        target_image = torch.from_numpy(style_img).to(device)
        target_image = target_image.permute(2, 0, 1)[None, ...]
        texture_name = url.split("/")[-1].split(".")[0]

        model_path = os.path.join(config['experiment_path'], f"{texture_name}")
        log_path = os.path.join(model_path, "logs")
        if not os.path.exists(model_path):
            os.makedirs(log_path)
        elif os.path.exists(os.path.join(model_path, "weights.pt")):
            print(f"A trained model for {texture_name} exists.")
            continue
        else:
            shutil.rmtree(model_path)
            os.makedirs(log_path)

        if wandb_enabled:
            name = config['experiment_name'] + f"-{texture_name}"
            wandb_run = wandb.init(project=config['wandb']['project'],
                                   name=name, dir=log_path, config=config)

        nca = get_nca_model(config, texture_name).to(device)

        # capturable=True only works with CUDA devices, so set based on device type
        capturable = device.type == 'cuda'
        opt = torch.optim.Adam(nca.parameters(), config['training']['lr'], capturable=capturable)

        lr_sched = None
        if 'type' not in config['training']['scheduler'] or config['training']['scheduler']['type'] == 'MultiStep':
            lr_sched = torch.optim.lr_scheduler.MultiStepLR(opt, **config['training']['scheduler']['attr'])
        elif config['training']['scheduler']['type'] == 'Cyclic':
            lr_sched = torch.optim.lr_scheduler.CyclicLR(opt, **config['training']['scheduler']['attr'])

        batch_size = config['training']['batch_size']
        iterations = config['training']['iterations']
        alpha = config['training']['overflow_weight']

        step_range = config['training']['nca']['step_range']
        inject_seed_step = config['training']['nca']['inject_seed_step']
        pool_size = config['training']['nca']['pool_size']

        with torch.no_grad():
            pool = nca.seed(pool_size).to(device)

        pbar = tqdm(range(iterations), desc=f"Training {idx + 1}/{len(image_paths)} on {texture_name}")
        for epoch in pbar:
            with torch.no_grad():
                batch_idx = np.random.choice(pool_size, batch_size, replace=False)
                x = pool[batch_idx]
                if epoch % inject_seed_step == 0:
                    x[:1] = nca.seed(1)

            step_n = np.random.randint(step_range[0], step_range[1])

            for _ in range(step_n):
                x = nca(x)

            overflow_loss = (x - x.clamp(-1.0, 1.0)).abs().sum()
            texture_loss, texture_loss_per_img = loss_fn(target_image, nca.to_rgb(x))
            loss = texture_loss + alpha * overflow_loss

            with torch.no_grad():
                loss.backward()
                for p in nca.parameters():
                    p.grad /= (p.grad.norm() + 1e-8)  # normalize gradients
                opt.step()
                opt.zero_grad()
                lr_sched.step()

                pool[batch_idx] = x

            if (epoch + 1) % config['training']['log_interval'] == 0:
                imgs = nca.to_rgb(x[:4]).permute([0, 2, 3, 1]).detach().cpu().numpy()
                imgs = np.hstack((np.clip(imgs, 0, 1) * 255.0).astype(np.uint8))
                if wandb_enabled:
                    wandb_run.log({'NCA Output': wandb.Image(imgs, caption='NCA Output')}, step=epoch)
                else:
                    Image.fromarray(imgs).save(f'{log_path}/epoch-{epoch}.png')

            if wandb_enabled:
                loss_log = {
                    'total': loss.item(),
                    'overflow': overflow_loss.item(),
                    'texture': texture_loss.item()
                }
                wandb_run.log(loss_log, step=epoch)

        torch.save(nca.state_dict(), f'{model_path}/weights.pt')

        if wandb_enabled:
            wandb_run.finish()
        del nca
        del opt


if __name__ == "__main__":
    args = parser.parse_args()
    with open(args.config, 'r') as stream:
        config = yaml.load(stream, Loader=yaml.FullLoader)

    config['data_dir'] = args.data_dir
    exp_name = config['experiment_name']
    exp_path = f'results_new/{exp_name}/'
    config['experiment_path'] = exp_path
    if not os.path.exists(exp_path):
        os.makedirs(exp_path)
    main(config)