first commit

.gitignore

@@ -0,0 +1,34 @@
+command
+pip_list
+
+__pycache__/
+*.py[cod]
+*$py.class
+.idea/
+
+stable-diffusion/CompVis
+stable-diffusion/src
+stable-diffusion/latent_diffusion.egg-info
+
+
+3DPortraitGAN_pyramid/models/*.pkl
+3DPortraitGAN_pyramid/models/*.ckpt
+3DPortraitGAN_pyramid/models/*.pt
+3DPortraitGAN_pyramid/models/*.json
+3DPortraitGAN_pyramid/out
+3DPortraitGAN_pyramid/smplx_models/smpl/*.pkl
+3DPortraitGAN_pyramid/training-runs
+
+stable-dreamfusion-3DPortrait/pretrained
+stable-dreamfusion-3DPortrait/output
+stable-dreamfusion-3DPortrait/smplx_models
+stable-dreamfusion-3DPortrait/transfer_data
+
+command.py
+temp.py
+
+*.pkl
+*.pth
+*.pt
+*.pth.tar
+./data_processing/data/J_regressor_extra.npy

3DPortraitGAN_pyramid/.gitignore

@@ -0,0 +1,8 @@
+models
+smplx_models
+transfer_data
+generate_inversion_results.py
+training/dataset-ref.py
+training-runs
+.vscode
+*.out

3DPortraitGAN_pyramid/README.md

@@ -0,0 +1,59 @@
+# 3DPortraitGAN_pyramid Training
+
+**Note: Upon the acceptance of our [3DPortraitGAN](https://arxiv.org/abs/2307.14770), we plan to release our 360°PHQ dataset to facilitate reproducibility of research. We encourage you to utilize our provided pre-trained models. Stay tuned for updates! **
+
+
+
+## Training
+
+```shell
+cd 3DPortraitGAN_pyramid
+
+# stage 1
+python  train.py \
+	--outdir=./training-runs/stage1 --cfg=full-head \
+	--data=$DATASET_PATH/360PHQ-512.zip  --seg_data=$DATASET_PATH/360PHQ-512-mask.zip \
+    --gpus=8 --batch=32  --gamma=5.0 --cbase=18432 --cmax=144 \
+	--gamma_seg=5.0 --use_torgb_raw=1 --decoder_activation="none" \
+	--bcg_reg_prob 0.2  --triplane_depth 3 --density_noise_fade_kimg 200 --density_reg 0 --back_repeat=1 \
+	--gen_pose_cond=True --gpc_reg_prob=0.7  --mirror=True  --data_rebalance=False  --image-snap=25  --kimg=20000 \
+	--neural_rendering_resolution_initial=64  \
+	--pose_loss_weight=10 --input_pose_params_reg_loss_weight=5 --input_pose_params_reg_loss_kimg=200 \
+	--train_g_pose_branch=True \
+	--explicitly_symmetry=True \
+	--metric_pose_sample_mode=G_predict 
+
+
+# stage 2
+python  train.py \
+	--outdir=./training-runs/stage2 --cfg=full-head \
+	--data=$DATASET_PATH/360PHQ-512.zip  --seg_data=$DATASET_PATH/360PHQ-512-mask.zip \
+	--gpus=8 --batch=32  --gamma=5.0 --cbase=18432 --cmax=144 \
+	--gamma_seg=5.0 --use_torgb_raw=1 --decoder_activation="none" \
+	--bcg_reg_prob 0.2  --triplane_depth 3 --density_noise_fade_kimg 200 --density_reg 0 --back_repeat=1 \
+	--gen_pose_cond=True --gpc_reg_prob=0.7  --mirror=True  --data_rebalance=False  --image-snap=25  --kimg=20000 \
+	--neural_rendering_resolution_initial=64  \
+	--pose_loss_weight=10 --input_pose_params_reg_loss_weight=5 --input_pose_params_reg_loss_kimg=200 \
+	--train_g_pose_branch=False \
+	--explicitly_symmetry=True \
+	--metric_pose_sample_mode=D_predict \
+	--resume=stage1.pkl  --resume_kimg=NUM_KIMGS
+
+# stage 3
+python  train.py \
+	--outdir=./training-runs/stage3 --cfg=full-head \
+	--data=$DATASET_PATH/360PHQ-512.zip  --seg_data=$DATASET_PATH/360PHQ-512-mask.zip \
+	--gpus=8 --batch=32  --gamma=5.0 --cbase=18432 --cmax=144 \
+	--gamma_seg=5.0 --use_torgb_raw=1 --decoder_activation="none" \
+	--bcg_reg_prob 0.2  --triplane_depth 3 --density_noise_fade_kimg 200 --density_reg 0 --back_repeat=1 \
+	--gen_pose_cond=True --gpc_reg_prob=0.7  --mirror=True  --data_rebalance=False  --image-snap=25  --kimg=20000 \
+	--neural_rendering_resolution_initial=64  --neural_rendering_resolution_final=128 \
+	--neural_rendering_resolution_fade_kimg=1000 \
+	--pose_loss_weight=10 --input_pose_params_reg_loss_weight=5 --input_pose_params_reg_loss_kimg=200 \
+	--train_g_pose_branch=False \
+	--explicitly_symmetry=True \
+	--metric_pose_sample_mode=D_predict \
+	--resume=stage2.pkl  --resume_kimg=NUM_KIMGS
+
+```
+

3DPortraitGAN_pyramid/calc_metrics.py

@@ -0,0 +1,232 @@
+# SPDX-FileCopyrightText: Copyright (c) 2021-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
+#
+# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
+# property and proprietary rights in and to this material, related
+# documentation and any modifications thereto. Any use, reproduction,
+# disclosure or distribution of this material and related documentation
+# without an express license agreement from NVIDIA CORPORATION or
+# its affiliates is strictly prohibited.
+
+"""Calculate quality metrics for previous training run or pretrained network pickle."""
+
+import os
+import click
+import json
+import tempfile
+import copy
+import torch
+
+import dnnlib
+import legacy
+from metrics import metric_main
+from metrics import metric_utils
+from torch_utils import training_stats
+from torch_utils import custom_ops
+from torch_utils import misc
+from torch_utils.ops import conv2d_gradfix
+
+#----------------------------------------------------------------------------
+
+def subprocess_fn(rank, args, temp_dir):
+    dnnlib.util.Logger(should_flush=True)
+
+    # Init torch.distributed.
+    if args.num_gpus > 1:
+        init_file = os.path.abspath(os.path.join(temp_dir, '.torch_distributed_init'))
+        if os.name == 'nt':
+            init_method = 'file:///' + init_file.replace('\\', '/')
+            torch.distributed.init_process_group(backend='gloo', init_method=init_method, rank=rank, world_size=args.num_gpus)
+        else:
+            init_method = f'file://{init_file}'
+            torch.distributed.init_process_group(backend='nccl', init_method=init_method, rank=rank, world_size=args.num_gpus)
+
+    # Init torch_utils.
+    sync_device = torch.device('cuda', rank) if args.num_gpus > 1 else None
+    training_stats.init_multiprocessing(rank=rank, sync_device=sync_device)
+    if rank != 0 or not args.verbose:
+        custom_ops.verbosity = 'none'
+
+    # Configure torch.
+    device = torch.device('cuda', rank)
+    torch.backends.cuda.matmul.allow_tf32 = False
+    torch.backends.cudnn.allow_tf32 = False
+    conv2d_gradfix.enabled = True
+
+    # Print network summary.
+    G = copy.deepcopy(args.G).eval().requires_grad_(False).to(device)
+    D = copy.deepcopy(args.D).eval().requires_grad_(False).to(device) if args.metric_pose_sample_mode == 'D_predict' else None
+
+
+    resample_filter = args.pose_predict_kwargs['resample_filter']
+    resample_filter = torch.tensor(resample_filter, device=device).to(torch.float32)
+
+    if rank == 0 and args.verbose:
+        z = torch.empty([1, G.z_dim], device=device)
+        c = torch.empty([1, G.c_dim], device=device)
+        misc.print_module_summary(G, [z, c])
+
+    # Calculate each metric.
+    for metric in args.metrics:
+        if rank == 0 and args.verbose:
+            print(f'Calculating {metric}...')
+        progress = metric_utils.ProgressMonitor(verbose=args.verbose)
+        # result_dict = metric_main.calc_metric(metric=metric, G=G, dataset_kwargs=args.dataset_kwargs,
+        #     num_gpus=args.num_gpus, rank=rank, device=device, progress=progress)
+
+
+        result_dict = metric_main.calc_metric(metric=metric,
+                                                      G=G,
+                                                      dataset_kwargs=args.dataset_kwargs,
+                                                      num_gpus=args.num_gpus,
+                                                      rank=rank,
+                                                      device=device,
+                                                      metric_pose_sample_mode = args.metric_pose_sample_mode,
+                                                      progress=progress,
+                                                      identical_c_p = args.identical_c_p,
+                                                      D = D,
+                                                      pose_predict_kwargs = {
+                                                         'neural_rendering_resolution':args.pose_predict_kwargs['neural_rendering_resolution'],
+                                'blur_sigma':args.pose_predict_kwargs['blur_sigma'],
+                                'resample_filter':resample_filter,
+                                'filter_mode':args.pose_predict_kwargs['filter_mode']
+                                                      }  if args.metric_pose_sample_mode == 'D_predict' else None
+                                                      )
+        if rank == 0:
+            metric_main.report_metric(result_dict, run_dir=args.run_dir, snapshot_pkl=args.network_pkl)
+        if rank == 0 and args.verbose:
+            print()
+
+    # Done.
+    if rank == 0 and args.verbose:
+        print('Exiting...')
+
+#----------------------------------------------------------------------------
+
+def parse_comma_separated_list(s):
+    if isinstance(s, list):
+        return s
+    if s is None or s.lower() == 'none' or s == '':
+        return []
+    return s.split(',')
+
+#----------------------------------------------------------------------------
+
+@click.command()
+@click.pass_context
+@click.option('network_pkl', '--network', help='Network pickle filename or URL', metavar='PATH', required=True)
+@click.option('pose_predict_kwargs', '--pose_predict_kwargs', help='Network pickle filename or URL', metavar='PATH', required=True)
+@click.option('--metric_pose_sample_mode', help='Type of metric_pose_sample ', metavar='STR',  type=click.Choice(['D_predict', 'G_predict']), required=False, default='G_predict')
+@click.option('--identical_c_p', help='Enable dataset x-flips  [default: look up]', type=bool, metavar='BOOL')
+
+@click.option('--metrics', help='Quality metrics', metavar='[NAME|A,B,C|none]', type=parse_comma_separated_list, default='fid50k_full', show_default=True)
+@click.option('--data', help='Dataset to evaluate against  [default: look up]', metavar='[ZIP|DIR]')
+@click.option('--seg_data', help='Dataset to evaluate against  [default: look up]', metavar='[ZIP|DIR]')
+@click.option('--mirror', help='Enable dataset x-flips  [default: look up]', type=bool, metavar='BOOL')
+@click.option('--gpus', help='Number of GPUs to use', type=int, default=1, metavar='INT', show_default=True)
+@click.option('--verbose', help='Print optional information', type=bool, default=True, metavar='BOOL', show_default=True)
+
+def calc_metrics(ctx, network_pkl, pose_predict_kwargs,metric_pose_sample_mode,identical_c_p ,metrics, data,seg_data, mirror, gpus, verbose):
+    """Calculate quality metrics for previous training run or pretrained network pickle.
+
+    Examples:
+
+    \b
+    # Previous training run: look up options automatically, save result to JSONL file.
+    python calc_metrics.py --metrics=eqt50k_int,eqr50k \\
+        --network=~/training-runs/00000-stylegan3-r-mydataset/network-snapshot-000000.pkl
+
+    \b
+    # Pre-trained network pickle: specify dataset explicitly, print result to stdout.
+    python calc_metrics.py --metrics=fid50k_full --data=~/datasets/ffhq-1024x1024.zip --mirror=1 \\
+        --network=https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-t-ffhq-1024x1024.pkl
+
+    \b
+    Recommended metrics:
+      fid50k_full  Frechet inception distance against the full dataset.
+      kid50k_full  Kernel inception distance against the full dataset.
+      pr50k3_full  Precision and recall againt the full dataset.
+      ppl2_wend    Perceptual path length in W, endpoints, full image.
+      eqt50k_int   Equivariance w.r.t. integer translation (EQ-T).
+      eqt50k_frac  Equivariance w.r.t. fractional translation (EQ-T_frac).
+      eqr50k       Equivariance w.r.t. rotation (EQ-R).
+
+    \b
+    Legacy metrics:
+      fid50k       Frechet inception distance against 50k real images.
+      kid50k       Kernel inception distance against 50k real images.
+      pr50k3       Precision and recall against 50k real images.
+      is50k        Inception score for CIFAR-10.
+    """
+    dnnlib.util.Logger(should_flush=True)
+
+    # Validate arguments.
+    args = dnnlib.EasyDict(metrics=metrics, num_gpus=gpus, network_pkl=network_pkl, verbose=verbose,metric_pose_sample_mode=metric_pose_sample_mode)
+    if not all(metric_main.is_valid_metric(metric) for metric in args.metrics):
+        ctx.fail('\n'.join(['--metrics can only contain the following values:'] + metric_main.list_valid_metrics()))
+    if not args.num_gpus >= 1:
+        ctx.fail('--gpus must be at least 1')
+
+    # Load network.
+    if not dnnlib.util.is_url(network_pkl, allow_file_urls=True) and not os.path.isfile(network_pkl):
+        ctx.fail('--network must point to a file or URL')
+    if args.verbose:
+        print(f'Loading network from "{network_pkl}"...')
+    with dnnlib.util.open_url(network_pkl, verbose=args.verbose) as f:
+        network_dict = legacy.load_network_pkl(f)
+        args.G = network_dict['G_ema'] # subclass of torch.nn.Module
+        args.D = network_dict['D_ema']
+
+    args.identical_c_p = identical_c_p
+
+    with open(pose_predict_kwargs, 'r') as f:
+        args.pose_predict_kwargs = json.load(f)
+
+
+    # Initialize dataset options.
+    if data is not None:
+        #args.dataset_kwargs = dnnlib.EasyDict(class_name='training.dataset.ImageFolderDataset', path=data)
+        args.dataset_kwargs = dnnlib.EasyDict(class_name='training.dataset.MaskLabeledDataset',
+                                         img_path=data,
+                                         seg_path = seg_data,
+                                         back_repeat =1,
+                                         use_labels=True, max_size=None, xflip=True,
+                                         data_rebalance=False,data_rebalance_idx_file=None)
+    elif network_dict['training_set_kwargs'] is not None:
+        args.dataset_kwargs = dnnlib.EasyDict(network_dict['training_set_kwargs'])
+    else:
+        ctx.fail('Could not look up dataset options; please specify --data')
+
+    # Finalize dataset options.
+    args.dataset_kwargs.resolution = args.G.img_resolution
+    args.dataset_kwargs.use_labels =True
+
+
+    # Print dataset options.
+    if args.verbose:
+        print('Dataset options:')
+        print(json.dumps(args.dataset_kwargs, indent=2))
+
+    # Locate run dir.
+    args.run_dir = None
+    if os.path.isfile(network_pkl):
+        pkl_dir = os.path.dirname(network_pkl)
+        if os.path.isfile(os.path.join(pkl_dir, 'training_options.json')):
+            args.run_dir = pkl_dir
+
+    # Launch processes.
+    if args.verbose:
+        print('Launching processes...')
+    torch.multiprocessing.set_start_method('spawn')
+    with tempfile.TemporaryDirectory() as temp_dir:
+        if args.num_gpus == 1:
+            subprocess_fn(rank=0, args=args, temp_dir=temp_dir)
+        else:
+            torch.multiprocessing.spawn(fn=subprocess_fn, args=(args, temp_dir), nprocs=args.num_gpus)
+
+#----------------------------------------------------------------------------
+
+if __name__ == "__main__":
+    calc_metrics() # pylint: disable=no-value-for-parameter
+
+#----------------------------------------------------------------------------