mirror of
https://github.com/Haoming02/sd-webui-old-photo-restoration.git
synced 2026-01-26 19:29:52 +00:00
470 lines
17 KiB
Python
470 lines
17 KiB
Python
# Copyright (c) Microsoft Corporation
|
|
|
|
import argparse
|
|
import torch
|
|
|
|
|
|
class BaseOptions:
|
|
def __init__(self):
|
|
self.parser = argparse.ArgumentParser()
|
|
self.initialized = False
|
|
|
|
def initialize(self):
|
|
# experiment specifics
|
|
self.parser.add_argument(
|
|
"--name",
|
|
type=str,
|
|
default="label2city",
|
|
help="name of the experiment. It decides where to store samples and models",
|
|
)
|
|
self.parser.add_argument(
|
|
"--gpu_ids",
|
|
type=str,
|
|
default="0",
|
|
help="gpu ids: e.g. 0 0,1,2, 0,2. use -1 for CPU",
|
|
)
|
|
self.parser.add_argument(
|
|
"--checkpoints_dir",
|
|
type=str,
|
|
default="./checkpoints",
|
|
help="models are saved here",
|
|
) ## note: to add this param when using philly
|
|
# self.parser.add_argument('--project_dir', type=str, default='./', help='the project is saved here') ################### This is necessary for philly
|
|
self.parser.add_argument(
|
|
"--outputs_dir", type=str, default="./outputs", help="models are saved here"
|
|
) ## note: to add this param when using philly Please end with '/'
|
|
self.parser.add_argument(
|
|
"--model", type=str, default="pix2pixHD", help="which model to use"
|
|
)
|
|
self.parser.add_argument(
|
|
"--norm",
|
|
type=str,
|
|
default="instance",
|
|
help="instance normalization or batch normalization",
|
|
)
|
|
self.parser.add_argument(
|
|
"--use_dropout", action="store_true", help="use dropout for the generator"
|
|
)
|
|
self.parser.add_argument(
|
|
"--data_type",
|
|
default=32,
|
|
type=int,
|
|
choices=[8, 16, 32],
|
|
help="Supported data type i.e. 8, 16, 32 bit",
|
|
)
|
|
self.parser.add_argument(
|
|
"--verbose", action="store_true", default=False, help="toggles verbose"
|
|
)
|
|
|
|
# input/output sizes
|
|
self.parser.add_argument(
|
|
"--batchSize", type=int, default=1, help="input batch size"
|
|
)
|
|
self.parser.add_argument(
|
|
"--loadSize", type=int, default=1024, help="scale images to this size"
|
|
)
|
|
self.parser.add_argument(
|
|
"--fineSize", type=int, default=512, help="then crop to this size"
|
|
)
|
|
self.parser.add_argument(
|
|
"--label_nc", type=int, default=35, help="# of input label channels"
|
|
)
|
|
self.parser.add_argument(
|
|
"--input_nc", type=int, default=3, help="# of input image channels"
|
|
)
|
|
self.parser.add_argument(
|
|
"--output_nc", type=int, default=3, help="# of output image channels"
|
|
)
|
|
|
|
# for setting inputs
|
|
self.parser.add_argument(
|
|
"--dataroot", type=str, default="./datasets/cityscapes/"
|
|
)
|
|
self.parser.add_argument(
|
|
"--resize_or_crop",
|
|
type=str,
|
|
default="scale_width",
|
|
help="scaling and cropping of images at load time [resize_and_crop|crop|scale_width|scale_width_and_crop]",
|
|
)
|
|
self.parser.add_argument(
|
|
"--serial_batches",
|
|
action="store_true",
|
|
help="if true, takes images in order to make batches, otherwise takes them randomly",
|
|
)
|
|
self.parser.add_argument(
|
|
"--no_flip",
|
|
action="store_true",
|
|
help="if specified, do not flip the images for data argumentation",
|
|
)
|
|
self.parser.add_argument(
|
|
"--nThreads", default=2, type=int, help="# threads for loading data"
|
|
)
|
|
self.parser.add_argument(
|
|
"--max_dataset_size",
|
|
type=int,
|
|
default=float("inf"),
|
|
help="Maximum number of samples allowed per dataset. If the dataset directory contains more than max_dataset_size, only a subset is loaded.",
|
|
)
|
|
|
|
# for displays
|
|
self.parser.add_argument(
|
|
"--display_winsize", type=int, default=512, help="display window size"
|
|
)
|
|
self.parser.add_argument(
|
|
"--tf_log",
|
|
action="store_true",
|
|
help="if specified, use tensorboard logging. Requires tensorflow installed",
|
|
)
|
|
|
|
# for generator
|
|
self.parser.add_argument(
|
|
"--netG", type=str, default="global", help="selects model to use for netG"
|
|
)
|
|
self.parser.add_argument(
|
|
"--ngf", type=int, default=64, help="# of gen filters in first conv layer"
|
|
)
|
|
self.parser.add_argument(
|
|
"--k_size", type=int, default=3, help="# kernel size conv layer"
|
|
)
|
|
self.parser.add_argument("--use_v2", action="store_true", help="use DCDCv2")
|
|
self.parser.add_argument("--mc", type=int, default=1024, help="# max channel")
|
|
self.parser.add_argument(
|
|
"--start_r", type=int, default=3, help="start layer to use resblock"
|
|
)
|
|
self.parser.add_argument(
|
|
"--n_downsample_global",
|
|
type=int,
|
|
default=4,
|
|
help="number of downsampling layers in netG",
|
|
)
|
|
self.parser.add_argument(
|
|
"--n_blocks_global",
|
|
type=int,
|
|
default=9,
|
|
help="number of residual blocks in the global generator network",
|
|
)
|
|
self.parser.add_argument(
|
|
"--n_blocks_local",
|
|
type=int,
|
|
default=3,
|
|
help="number of residual blocks in the local enhancer network",
|
|
)
|
|
self.parser.add_argument(
|
|
"--n_local_enhancers",
|
|
type=int,
|
|
default=1,
|
|
help="number of local enhancers to use",
|
|
)
|
|
self.parser.add_argument(
|
|
"--niter_fix_global",
|
|
type=int,
|
|
default=0,
|
|
help="number of epochs that we only train the outmost local enhancer",
|
|
)
|
|
|
|
self.parser.add_argument(
|
|
"--load_pretrain",
|
|
type=str,
|
|
default="",
|
|
help="load the pretrained model from the specified location",
|
|
)
|
|
|
|
# for instance-wise features
|
|
self.parser.add_argument(
|
|
"--no_instance",
|
|
action="store_true",
|
|
help="if specified, do *not* add instance map as input",
|
|
)
|
|
self.parser.add_argument(
|
|
"--instance_feat",
|
|
action="store_true",
|
|
help="if specified, add encoded instance features as input",
|
|
)
|
|
self.parser.add_argument(
|
|
"--label_feat",
|
|
action="store_true",
|
|
help="if specified, add encoded label features as input",
|
|
)
|
|
self.parser.add_argument(
|
|
"--feat_num", type=int, default=3, help="vector length for encoded features"
|
|
)
|
|
self.parser.add_argument(
|
|
"--load_features",
|
|
action="store_true",
|
|
help="if specified, load precomputed feature maps",
|
|
)
|
|
self.parser.add_argument(
|
|
"--n_downsample_E",
|
|
type=int,
|
|
default=4,
|
|
help="# of downsampling layers in encoder",
|
|
)
|
|
self.parser.add_argument(
|
|
"--nef",
|
|
type=int,
|
|
default=16,
|
|
help="# of encoder filters in the first conv layer",
|
|
)
|
|
self.parser.add_argument(
|
|
"--n_clusters", type=int, default=10, help="number of clusters for features"
|
|
)
|
|
|
|
# diy
|
|
self.parser.add_argument(
|
|
"--self_gen", action="store_true", help="self generate"
|
|
)
|
|
self.parser.add_argument(
|
|
"--mapping_n_block",
|
|
type=int,
|
|
default=3,
|
|
help="number of resblock in mapping",
|
|
)
|
|
self.parser.add_argument(
|
|
"--map_mc", type=int, default=64, help="max channel of mapping"
|
|
)
|
|
self.parser.add_argument("--kl", type=float, default=0, help="KL Loss")
|
|
self.parser.add_argument(
|
|
"--load_pretrainA",
|
|
type=str,
|
|
default="",
|
|
help="load the pretrained model from the specified location",
|
|
)
|
|
self.parser.add_argument(
|
|
"--load_pretrainB",
|
|
type=str,
|
|
default="",
|
|
help="load the pretrained model from the specified location",
|
|
)
|
|
self.parser.add_argument("--feat_gan", action="store_true")
|
|
self.parser.add_argument("--no_cgan", action="store_true")
|
|
self.parser.add_argument("--map_unet", action="store_true")
|
|
self.parser.add_argument("--map_densenet", action="store_true")
|
|
self.parser.add_argument("--fcn", action="store_true")
|
|
self.parser.add_argument(
|
|
"--is_image", action="store_true", help="train image recon only pair data"
|
|
)
|
|
self.parser.add_argument("--label_unpair", action="store_true")
|
|
self.parser.add_argument("--mapping_unpair", action="store_true")
|
|
self.parser.add_argument("--unpair_w", type=float, default=1.0)
|
|
self.parser.add_argument("--pair_num", type=int, default=-1)
|
|
self.parser.add_argument("--Gan_w", type=float, default=1)
|
|
self.parser.add_argument("--feat_dim", type=int, default=-1)
|
|
self.parser.add_argument("--abalation_vae_len", type=int, default=-1)
|
|
|
|
######################### useless, just to cooperate with docker
|
|
self.parser.add_argument("--gpu", type=str)
|
|
self.parser.add_argument("--dataDir", type=str)
|
|
self.parser.add_argument("--modelDir", type=str)
|
|
self.parser.add_argument("--logDir", type=str)
|
|
self.parser.add_argument("--data_dir", type=str)
|
|
|
|
self.parser.add_argument("--use_skip_model", action="store_true")
|
|
self.parser.add_argument("--use_segmentation_model", action="store_true")
|
|
|
|
self.parser.add_argument("--spatio_size", type=int, default=64)
|
|
self.parser.add_argument("--test_random_crop", action="store_true")
|
|
##########################
|
|
|
|
self.parser.add_argument("--contain_scratch_L", action="store_true")
|
|
self.parser.add_argument(
|
|
"--mask_dilation", type=int, default=0
|
|
) ## Don't change the input, only dilation the mask
|
|
|
|
self.parser.add_argument(
|
|
"--irregular_mask",
|
|
type=str,
|
|
default="",
|
|
help="This is the root of the mask",
|
|
)
|
|
self.parser.add_argument(
|
|
"--mapping_net_dilation",
|
|
type=int,
|
|
default=1,
|
|
help="This parameter is the dilation size of the translation net",
|
|
)
|
|
|
|
self.parser.add_argument(
|
|
"--VOC",
|
|
type=str,
|
|
default="VOC_RGB_JPEGImages.bigfile",
|
|
help="The root of VOC dataset",
|
|
)
|
|
|
|
self.parser.add_argument(
|
|
"--non_local", type=str, default="", help="which non_local setting"
|
|
)
|
|
self.parser.add_argument(
|
|
"--NL_fusion_method",
|
|
type=str,
|
|
default="add",
|
|
help="how to fuse the origin feature and nl feature",
|
|
)
|
|
self.parser.add_argument(
|
|
"--NL_use_mask",
|
|
action="store_true",
|
|
help="If use mask while using Non-local mapping model",
|
|
)
|
|
self.parser.add_argument(
|
|
"--correlation_renormalize",
|
|
action="store_true",
|
|
help="Since after mask out the correlation matrix(which is softmaxed), the sum is not 1 any more, enable this param to re-weight",
|
|
)
|
|
|
|
self.parser.add_argument(
|
|
"--Smooth_L1", action="store_true", help="Use L1 Loss in image level"
|
|
)
|
|
|
|
self.parser.add_argument(
|
|
"--face_restore_setting",
|
|
type=int,
|
|
default=1,
|
|
help="This is for the aligned face restoration",
|
|
)
|
|
self.parser.add_argument("--face_clean_url", type=str, default="")
|
|
self.parser.add_argument("--syn_input_url", type=str, default="")
|
|
self.parser.add_argument("--syn_gt_url", type=str, default="")
|
|
|
|
self.parser.add_argument(
|
|
"--test_on_synthetic",
|
|
action="store_true",
|
|
help="If you want to test on the synthetic data, enable this parameter",
|
|
)
|
|
|
|
self.parser.add_argument(
|
|
"--use_SN", action="store_true", help="Add SN to every parametric layer"
|
|
)
|
|
|
|
self.parser.add_argument(
|
|
"--use_two_stage_mapping",
|
|
action="store_true",
|
|
help="choose the model which uses two stage",
|
|
)
|
|
|
|
self.parser.add_argument("--L1_weight", type=float, default=10.0)
|
|
self.parser.add_argument("--softmax_temperature", type=float, default=1.0)
|
|
self.parser.add_argument(
|
|
"--patch_similarity",
|
|
action="store_true",
|
|
help="Enable this denotes using 3*3 patch to calculate similarity",
|
|
)
|
|
self.parser.add_argument(
|
|
"--use_self",
|
|
action="store_true",
|
|
help="Enable this denotes that while constructing the new feature maps, using original feature (diagonal == 1)",
|
|
)
|
|
|
|
self.parser.add_argument("--use_own_dataset", action="store_true")
|
|
|
|
self.parser.add_argument(
|
|
"--test_hole_two_folders",
|
|
action="store_true",
|
|
help="Enable this parameter means test the restoration with inpainting given twp folders which are mask and old respectively",
|
|
)
|
|
|
|
self.parser.add_argument(
|
|
"--no_hole",
|
|
action="store_true",
|
|
help="While test the full_model on non_scratch data, do not add random mask into the real old photos",
|
|
) ## Only for testing
|
|
self.parser.add_argument(
|
|
"--random_hole",
|
|
action="store_true",
|
|
help="While training the full model, 50% probability add hole",
|
|
)
|
|
|
|
self.parser.add_argument(
|
|
"--NL_res", action="store_true", help="NL+Resdual Block"
|
|
)
|
|
|
|
self.parser.add_argument(
|
|
"--image_L1", action="store_true", help="Image level loss: L1"
|
|
)
|
|
self.parser.add_argument(
|
|
"--hole_image_no_mask",
|
|
action="store_true",
|
|
help="while testing, give hole image but not give the mask",
|
|
)
|
|
|
|
self.parser.add_argument(
|
|
"--down_sample_degradation",
|
|
action="store_true",
|
|
help="down_sample the image only, corresponds to [down_sample_face]",
|
|
)
|
|
|
|
self.parser.add_argument(
|
|
"--norm_G",
|
|
type=str,
|
|
default="spectralinstance",
|
|
help="The norm type of Generator",
|
|
)
|
|
self.parser.add_argument(
|
|
"--init_G",
|
|
type=str,
|
|
default="xavier",
|
|
help="normal|xavier|xavier_uniform|kaiming|orthogonal|none",
|
|
)
|
|
|
|
self.parser.add_argument("--use_new_G", action="store_true")
|
|
self.parser.add_argument("--use_new_D", action="store_true")
|
|
|
|
self.parser.add_argument(
|
|
"--only_voc",
|
|
action="store_true",
|
|
help="test the trianed celebA face model using VOC face",
|
|
)
|
|
|
|
self.parser.add_argument(
|
|
"--cosin_similarity",
|
|
action="store_true",
|
|
help="For non-local, using cosin to calculate the similarity",
|
|
)
|
|
|
|
self.parser.add_argument(
|
|
"--downsample_mode",
|
|
type=str,
|
|
default="nearest",
|
|
help="For partial non-local, choose how to downsample the mask",
|
|
)
|
|
|
|
self.parser.add_argument(
|
|
"--mapping_exp",
|
|
type=int,
|
|
default=0,
|
|
help="Default 0: original PNL|1: Multi-Scale Patch Attention",
|
|
)
|
|
self.parser.add_argument(
|
|
"--inference_optimize", action="store_true", help="optimize the memory cost"
|
|
)
|
|
|
|
self.initialized = True
|
|
|
|
def parse(self, custom_args: list):
|
|
|
|
assert len(custom_args) > 0, "Manually Pass Arguments!"
|
|
|
|
if not self.initialized:
|
|
self.initialize()
|
|
|
|
self.opt = self.parser.parse_args(custom_args)
|
|
self.opt.isTrain = self.isTrain # train or test
|
|
|
|
str_ids = self.opt.gpu_ids.split(",")
|
|
self.opt.gpu_ids = []
|
|
for str_id in str_ids:
|
|
int_id = int(str_id)
|
|
if int_id >= 0:
|
|
self.opt.gpu_ids.append(int_id)
|
|
|
|
# set gpu ids
|
|
if torch.cuda.is_available() and len(self.opt.gpu_ids) > 0:
|
|
if torch.cuda.device_count() > self.opt.gpu_ids[0]:
|
|
try:
|
|
torch.cuda.set_device(self.opt.gpu_ids[0])
|
|
except:
|
|
print("Failed to set GPU device. Using CPU...")
|
|
|
|
else:
|
|
print("Invalid GPU ID. Using CPU...")
|
|
|
|
return self.opt
|