main.py

import argparse
import traceback
import logging
import yaml
import sys
import os
import torch
import numpy as np
import json

from styleremoval import StyleRemoval
from styleremoval_image import StyleRemovalImage
from styleremoval_gaussian import StyleRemovalGaussian
from styleremoval_image_gaussian import StyleRemovalImageGaussian


from styletransfer import StyleTransfer
from configs.paths_config import HYBRID_MODEL_PATHS

def parse_args_and_config():
    parser = argparse.ArgumentParser(description=globals()['__doc__'])

    # Mode

    # Default
    parser.add_argument('--config', type=str, default='imagenet.yml', help='Path to the config file')
    parser.add_argument('--seed', type=int, default=1234, help='Random seed')
    parser.add_argument('--exp', type=str, default='./runs_gaussian/test', help='Path for saving running related data.')
    parser.add_argument('--comment', type=str, default='', help='A string for experiment comment')
    parser.add_argument('--verbose', type=str, default='info', help='Verbose level: info | debug | warning | critical')
    parser.add_argument('--ni', type=int, default=1,  help="No interaction. Suitable for Slurm Job launcher")
    parser.add_argument('--align_face', type=int, default=1, help='align face or not')

    # Image
    parser.add_argument('--style_image', type=str, default='gogh.jpg' , help='Style image')
    parser.add_argument('--removal_mode', type=str, default='gaussian' , help='Removal mode')
    parser.add_argument('--image_size', type=int, default=512, help='Image Size')
    parser.add_argument('--content_images', type=str, default='imagenet_subset' , help='Content folder')

    # Sampling
    parser.add_argument('--gaussian_kernel', type=float, default=1.5 , help='gaussian kernel for gaussian style removal')
    parser.add_argument('--t_0_remove', type=int, default=603, help='Return step in [0, 1000)')
    parser.add_argument('--t_0_transfer', type=int, default=601, help='Return step in [0, 1000)')
    parser.add_argument('--k_r', type=int, default=50, help='Return step in [0, 1000)')
    parser.add_argument('--n_inv_step', type=int, default=40, help='# of steps during generative pross for inversion')
    parser.add_argument('--n_train_step', type=int, default=6, help='# of steps during generative pross for train')
    parser.add_argument('--n_test_step', type=int, default=40, help='# of steps during generative pross for test')
    parser.add_argument('--sample_type', type=str, default='ddim', help='ddpm for Markovian sampling, ddim for non-Markovian sampling')
    parser.add_argument('--eta', type=float, default=0.0, help='Controls of varaince of the generative process')

    # Train & Test
    parser.add_argument('--do_train', type=int, default=1, help='Whether to train or not during CLIP finetuning')
    parser.add_argument('--do_test', type=int, default=1, help='Whether to test or not during CLIP finetuning')
    parser.add_argument('--save_train_image', type=int, default=1, help='Wheter to save training results during CLIP fineuning')
    parser.add_argument('--bs_train', type=int, default=1, help='Training batch size during CLIP fineuning')
    parser.add_argument('--bs_test', type=int, default=1, help='Test batch size during CLIP fineuning')
    parser.add_argument('--n_precomp_img', type=int, default=50, help='# of images to precompute latents')
    parser.add_argument('--n_train_img', type=int, default=50, help='# of training images')
    parser.add_argument('--n_test_img', type=int, default=5, help='# of test images')
    parser.add_argument('--deterministic_inv', type=int, default=1, help='Whether to use deterministic inversion during inference')
    parser.add_argument('--hybrid_noise', type=int, default=0, help='Whether to change multiple attributes by mixing multiple models')
    parser.add_argument('--model_ratio', type=float, default=1, help='Degree of change, noise ratio from original and finetuned model.')


    # Loss & Optimization
    parser.add_argument('--dir_loss', type=float, default=1., help='Weights of direction loss')
    parser.add_argument('--l1_loss_w', type=float, default=10., help='Weights of L1 loss')
    parser.add_argument('--style_loss_w', type=float, default=1., help='Weights of style loss')
    parser.add_argument('--clip_model_name', type=str, default='ViT-B/16', help='ViT-B/16, ViT-B/32, RN50x16 etc')
    parser.add_argument('--lr_clip_finetune', type=float, default=4e-6, help='Initial learning rate for finetuning')
    parser.add_argument('--lr_clip_lat_opt', type=float, default=2e-2, help='Initial learning rate for latent optim')
    parser.add_argument('--n_iter', type=int, default=5, help='# of iterations of a generative process with `n_train_img` images')
    parser.add_argument('--scheduler', type=int, default=1, help='Whether to increase the learning rate')
    parser.add_argument('--sch_gamma', type=float, default=1.2, help='Scheduler gamma')

    args = parser.parse_args()

    # parse config file
    with open(os.path.join('configs', args.config), 'r') as f:
        config = yaml.safe_load(f)
    new_config = dict2namespace(config)

    image_name = args.style_image.split('.')[0]
    args.exp = args.exp + f'_FT_{new_config.data.category}_{image_name}_s{args.image_size}_t{args.t_0_remove}_ninv{args.n_inv_step}_ngen{args.n_train_step}_dir_{args.dir_loss}_l1_{args.l1_loss_w}_st_{args.style_loss_w}_lr_{args.lr_clip_finetune}_{args.removal_mode}'

    level = getattr(logging, args.verbose.upper(), None)
    if not isinstance(level, int):
        raise ValueError('level {} not supported'.format(args.verbose))

    handler1 = logging.StreamHandler()
    formatter = logging.Formatter('%(levelname)s - %(filename)s - %(asctime)s - %(message)s')
    handler1.setFormatter(formatter)
    logger = logging.getLogger()
    logger.addHandler(handler1)
    logger.setLevel(level)

    os.makedirs(args.exp, exist_ok=True)
    os.makedirs('checkpoint', exist_ok=True)
    os.makedirs('precomputed', exist_ok=True)
    os.makedirs('runs', exist_ok=True)
    os.makedirs(args.exp, exist_ok=True)
    args.image_folder = os.path.join(args.exp, 'image_samples')
    if not os.path.exists(args.image_folder):
        os.makedirs(args.image_folder)
    else:
        overwrite = False
        if args.ni:
            overwrite = True
        else:
            response = input("Image folder already exists. Overwrite? (Y/N)")
            if response.upper() == 'Y':
                overwrite = True

        if overwrite:
            # shutil.rmtree(args.image_folder)
            os.makedirs(args.image_folder, exist_ok=True)
        else:
            print("Output image folder exists. Program halted.")
            sys.exit(0)

    # add device
    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
    logging.info("Using device: {}".format(device))
    new_config.device = device

    # set random seed
    torch.manual_seed(args.seed)
    np.random.seed(args.seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(args.seed)

    torch.backends.cudnn.benchmark = True

    return args, new_config


def dict2namespace(config):
    namespace = argparse.Namespace()
    for key, value in config.items():
        if isinstance(value, dict):
            new_value = dict2namespace(value)
        else:
            new_value = value
        setattr(namespace, key, new_value)
    return namespace


def main():
    args, config = parse_args_and_config()
    print(">" * 80)
    logging.info("Exp instance id = {}".format(os.getpid()))
    logging.info("Exp comment = {}".format(args.comment))
    logging.info("Config =")
    print("<" * 80)

    assert args.removal_mode in {'diffusion', 'gaussian'}
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)

    start.record()

    exists_tr = True

    for mode in ['train', 'test']:
        pairs_path = os.path.join('precomputed/',
                                          f'{config.data.category}_{mode}_t{args.t_0_remove}_size{args.image_size}_nim{args.n_precomp_img}_ninv{args.n_inv_step}_{args.removal_mode}_pairs.pth')
        exists_tr = exists_tr and os.path.exists(pairs_path)
    
    image_name = args.style_image.split('.')[0]
    style_pairs_path = os.path.join('precomputed/',
                                          f'{config.data.category}_style_{image_name}_t{args.t_0_remove}_size{args.image_size}_nim{args.n_precomp_img}_ninv{args.n_inv_step}_{args.removal_mode}_pairs.pth')
    exists_style =  os.path.exists(style_pairs_path)
    
    if not exists_tr:
        if args.removal_mode == 'diffusion':
            w = StyleRemoval(args, config)
            w.remove_style()
        elif args.removal_mode == 'gaussian':
            w = StyleRemovalGaussian(args, config)
            w.remove_style()
    
    if not exists_style:
        if args.removal_mode == 'diffusion':
            w = StyleRemovalImage(args, config)
            w.remove_style()
        elif args.removal_mode == 'gaussian':
            w = StyleRemovalImageGaussian(args, config)
            w.remove_style()

    w = StyleTransfer(args, config)
    w.transfer_style()

    end.record()
    torch.cuda.synchronize()

    train_time = start.elapsed_time(end)
    time_dict = {'time': train_time}
    with open(os.path.join(args.image_folder, 'time.json'), 'w') as f:
        json.dump(time_dict, f)



    return 0


if __name__ == '__main__':
    sys.exit(main())