Align Generator selection (#3); Fix doc's typos.

Author: Xiaoming Zhao

docs/TRAIN_EVAL.md

@@ -125,11 +125,11 @@ If everyting goes well, you should observe the following folder structure:
 ```
 .
 +-- ckpts
-| +-- afhqcat.pkl # file
-| +-- metfaces.pkl # file
 | +-- stylegan2_pretrained # folder
+| | +-- afhqcat.pkl # file
+| | +-- metfaces.pkl # file
 | | +-- transfer-learning-source-nets # folder
-+-- dataset
++-- runtime_dataset
 | +-- ffhq256x256.zip # file
 | +-- ffhq256_deep3dface_coeffs # folder
 | +-- ffhq512x512.zip # file
@@ -179,7 +179,7 @@ The command to evaluate the trained model is in [eval.sh](../gmpi/eval/eval.sh).
 
 Run the following command to evalute the model:
 ```bash
-bash ${GMPI_ROOT}/eval/eval.sh \
+bash ${GMPI_ROOT}/gmpi/eval/eval.sh \
  ${GMPI_ROOT} \
  FFHQ512 \  # this can be FFHQ256, FFHQ512, FFHQ1024, AFHQCat, or MetFaces
  exp_id \  # this is your experiment ID

gmpi/eval/common.py

@@ -22,32 +22,38 @@ def setup_model(opt, config, metadata, mpi_xyz_input, mpi_xyz_only_z, vis_mesh=F
  n_g_out_channels = 4
  n_g_out_planes = opt.nplanes
 
- if "depth2alpha_n_z_bins" in config.GMPI.MPI and config.GMPI.MPI.depth2alpha_n_z_bins is not None:
- from gmpi.models.networks.networks_vanilla_depth2alpha import Generator as StyleGAN2Generator
-
- if config.GMPI.TRAIN.normalized_xyz_range == "01":
- depth2alpha_z_range = 1.0
- elif config.GMPI.TRAIN.normalized_xyz_range == "-11":
- depth2alpha_z_range = 2.0
- else:
- raise ValueError
- else:
+ if config.GMPI.TRAIN.normalized_xyz_range == "01":
  depth2alpha_z_range = 1.0
-
- if "depth2alpha_n_z_bins" not in config.GMPI.MPI:
- config.defrost()
- config.GMPI.MPI.depth2alpha_n_z_bins = None
- config.freeze()
-
- if config.GMPI.MODEL.STYLEGAN2.torgba_cond_on_pos_enc != "none":
+ elif config.GMPI.TRAIN.normalized_xyz_range == "-11":
+ depth2alpha_z_range = 2.0
+ else:
+ raise ValueError
+
+ if "depth2alpha_n_z_bins" not in config.GMPI.MPI:
+ config.defrost()
+ config.GMPI.MPI.depth2alpha_n_z_bins = None
+ config.freeze()
+
+ if config.GMPI.MODEL.STYLEGAN2.torgba_cond_on_pos_enc != "none":
+ if config.GMPI.MODEL.STYLEGAN2.torgba_cond_on_pos_enc == "depth2alpha":
+ print("\nGenerator comes from depth2alpha\n")
+ from gmpi.models.networks.networks_vanilla_depth2alpha import Generator as StyleGAN2Generator
+ elif config.GMPI.MODEL.STYLEGAN2.torgba_cond_on_pos_enc == "normalize_add_z":
  if config.GMPI.MODEL.STYLEGAN2.torgba_cond_on_pos_enc_embed_func in ["learnable_param"]:
+ print("\nGenerator comes from learnable_param\n")
+ n_g_out_planes = config.GMPI.MPI.n_gen_planes
  from gmpi.models.networks.networks_pos_enc_learnable_param import Generator as StyleGAN2Generator
-
- n_g_out_planes = n_g_out_planes = config.GMPI.MPI.n_gen_planes
- else:
+ elif config.GMPI.MODEL.STYLEGAN2.torgba_cond_on_pos_enc_embed_func in ["modulated_lrelu"]:
+ print("\nGenerator comes from cond_on_depth\n")
  from gmpi.models.networks.networks_cond_on_pos_enc import Generator as StyleGAN2Generator
+ else:
+ raise NotImplementedError
  else:
- from gmpi.models.networks.networks_vanilla import Generator as StyleGAN2Generator
+ raise NotImplementedError
+ else:
+ print("\nGenerator comes from vanilla\n")
+ n_g_out_planes = config.GMPI.MPI.n_gen_planes
+ from gmpi.models.networks.networks_vanilla import Generator as StyleGAN2Generator
 
  synthesis_kwargs = convert_cfg_to_dict(config.GMPI.MODEL.STYLEGAN2.synthesis_kwargs)
  synthesis_kwargs_D = convert_cfg_to_dict(config.GMPI.MODEL.STYLEGAN2.synthesis_kwargs)

gmpi/eval/prepare_fake_data.py

@@ -106,9 +106,15 @@ def main(opt):
  # NOTE: Deep3DFaceRecon can only provide mask and depth of 224x224
  # https://github.com/sicxu/Deep3DFaceRecon_pytorch
  metadata["img_size"] = 224
+ 
+ if config.GMPI.MODEL.STYLEGAN2.torgba_cond_on_pos_enc == "none":
+ # Vanilla version. The number of planes must be same in training and evaluation.
+ n_mpi_planes = config.GMPI.MPI.n_gen_planes
+ else:
+ n_mpi_planes = opt.nplanes
 
  mpi_renderer = MPIRenderer(
- n_mpi_planes=opt.nplanes, # config.GMPI.MPI.n_gen_planes,
+ n_mpi_planes=n_mpi_planes, # config.GMPI.MPI.n_gen_planes,
  plane_min_d=metadata["ray_start"],
  plane_max_d=metadata["ray_end"],
  plan_spatial_enlarge_factor=config.GMPI.MPI.CAM_SETUP.spatial_enlarge_factor,

gmpi/models/networks/networks_cond_on_pos_enc.py

@@ -1234,7 +1234,7 @@ def __init__(self,
  gen_alpha_largest_res = 256,
  G_final_img_act = "none",
  depth2alpha_z_range=1.0,
- depth2alpha_n_z_bins=10, 
+ depth2alpha_n_z_bins=None, 
  ):
  super().__init__()
 

gmpi/models/networks/networks_vanilla.py

@@ -604,7 +604,7 @@ def __init__(self,
  self.conv_clamp = synthesis_kwargs["conv_clamp"]
 
  def forward(self, z, c, mpi_xyz_coords, xyz_coords_only_z, n_planes, truncation_psi=1, truncation_cutoff=None,
- enable_mapping_grad=True, enable_syn_feat_net_grad=True, **synthesis_kwargs):
+ enable_mapping_grad=True, enable_syn_feat_net_grad=True, z_interpolation_ws=None, **synthesis_kwargs):
 
  ws = self.mapping(z, c, truncation_psi=truncation_psi, truncation_cutoff=truncation_cutoff)
  img = self.synthesis(ws, **synthesis_kwargs)

gmpi/train.py

@@ -86,12 +86,21 @@ def train(rank, world_size, config, master_port, run_dataset):
  # fmt: off
  if config.GMPI.MODEL.STYLEGAN2.torgba_cond_on_pos_enc != "none":
  if config.GMPI.MODEL.STYLEGAN2.torgba_cond_on_pos_enc == "depth2alpha":
+ print("\nGenerator comes from depth2alpha\n")
  from gmpi.models.networks.networks_vanilla_depth2alpha import Generator as StyleGAN2Generator
- elif config.GMPI.MODEL.STYLEGAN2.torgba_cond_on_pos_enc_embed_func in ["learnable_param"]:
- from gmpi.models.networks.networks_pos_enc_learnable_param import Generator as StyleGAN2Generator
+ elif config.GMPI.MODEL.STYLEGAN2.torgba_cond_on_pos_enc == "normalize_add_z":
+ if config.GMPI.MODEL.STYLEGAN2.torgba_cond_on_pos_enc_embed_func in ["learnable_param"]:
+ print("\nGenerator comes from learnable_param\n")
+ from gmpi.models.networks.networks_pos_enc_learnable_param import Generator as StyleGAN2Generator
+ elif config.GMPI.MODEL.STYLEGAN2.torgba_cond_on_pos_enc_embed_func in ["modulated_lrelu"]:
+ print("\nGenerator comes from cond_on_depth\n")
+ from gmpi.models.networks.networks_cond_on_pos_enc import Generator as StyleGAN2Generator
+ else:
+ raise NotImplementedError
  else:
- from gmpi.models.networks.networks_cond_on_pos_enc import Generator as StyleGAN2Generator
+ raise NotImplementedError
  else:
+ print("\nGenerator comes from vanilla\n")
  from gmpi.models.networks.networks_vanilla import Generator as StyleGAN2Generator
  # fmt: on