-
Notifications
You must be signed in to change notification settings - Fork 90
/
main.py
112 lines (92 loc) · 3.8 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
import argparse
import traceback
import shutil
import logging
import yaml
import sys
import os
import torch
import numpy as np
import torch.utils.tensorboard as tb
import copy
from runners.image_editing import Diffusion
def parse_args_and_config():
parser = argparse.ArgumentParser(description=globals()['__doc__'])
parser.add_argument('--config', type=str, required=True, help='Path to the config file')
parser.add_argument('--seed', type=int, default=1234, help='Random seed')
parser.add_argument('--exp', type=str, default='exp', 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('--sample', action='store_true', help='Whether to produce samples from the model')
parser.add_argument('-i', '--image_folder', type=str, default='images', help="The folder name of samples")
parser.add_argument('--ni', action='store_true', help="No interaction. Suitable for Slurm Job launcher")
parser.add_argument('--npy_name', type=str, required=True)
parser.add_argument('--sample_step', type=int, default=3, help='Total sampling steps')
parser.add_argument('--t', type=int, default=400, help='Sampling noise scale')
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)
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(os.path.join(args.exp, 'image_samples'), exist_ok=True)
args.image_folder = os.path.join(args.exp, 'image_samples', args.image_folder)
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)
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)
try:
runner = Diffusion(args, config)
runner.image_editing_sample()
except Exception:
logging.error(traceback.format_exc())
return 0
if __name__ == '__main__':
sys.exit(main())