add code from the original repo

.gitignore

@@ -0,0 +1,25 @@
+__pycache__/
+.cache/
+datasets/
+outputs/
+out/
+out2/
+debug/
+checkpoint/
+*.zip
+*.npy
+core
+history/
+tools/*
+tools
+eval_outputs/
+pretrained/
+.nfs00d20000091cb3390001ead3
+scripts/research/
+
+.idea
+.vscode
+.github
+.ipynb_checkpoints/
+_screenshots/
+flagged

README.md

@@ -1,37 +1,75 @@
----
-title: StyleNeRF
-emoji: 🏃
-colorFrom: blue
-colorTo: indigo
-sdk: gradio
-app_file: app.py
-pinned: false
----
+# StyleNeRF: A Style-based 3D-Aware Generator for High-resolution Image Synthesis</sub>
 
-# Configuration
+![Random Sample](./docs/random_sample.jpg)
 
-`title`: _string_  
-Display title for the Space
+**StyleNeRF: A Style-based 3D-Aware Generator for High-resolution Image Synthesis**<br>
+Jiatao Gu, Lingjie Liu, Peng Wang, Christian Theobalt<br>
+### [Project Page](http://jiataogu.me/style_nerf) | [Video](http://jiataogu.me/style_nerf) | [Paper](https://arxiv.org/abs/2110.08985) | [Data](#dataset)<br>
 
-`emoji`: _string_  
-Space emoji (emoji-only character allowed)
+Abstract: *We propose StyleNeRF, a 3D-aware generative model for photo-realistic high-resolution image synthesis with high multi-view consistency, which can be trained on unstructured 2D images. Existing approaches either cannot synthesize high-resolution images with fine details or yield noticeable 3D-inconsistent artifacts. In addition, many of them lack control over style attributes and explicit 3D camera poses. StyleNeRF integrates the neural radiance field (NeRF) into a style-based generator to tackle the aforementioned challenges, i.e., improving rendering efficiency and 3D consistency for high-resolution image generation. We perform volume rendering only to produce a low-resolution feature map and progressively apply upsampling in 2D to address the first issue. To mitigate the inconsistencies caused by 2D upsampling, we propose multiple designs, including a better upsampler and a new regularization loss. With these designs, StyleNeRF can synthesize high-resolution images at interactive rates while preserving 3D consistency at high quality. StyleNeRF also enables control of camera poses and different levels of styles, which can generalize to unseen views. It also supports challenging tasks, including zoom-in and-out, style mixing, inversion, and semantic editing.*
 
-`colorFrom`: _string_  
-Color for Thumbnail gradient (red, yellow, green, blue, indigo, purple, pink, gray)
+## Requirements
+The codebase is tested on 
+* Python 3.7
+* PyTorch 1.7.1
+* 8 Nvidia GPU (Tesla V100 32GB) with CUDA version 11.0
 
-`colorTo`: _string_  
-Color for Thumbnail gradient (red, yellow, green, blue, indigo, purple, pink, gray)
+For additional python libraries, please install by:
 
-`sdk`: _string_  
-Can be either `gradio` or `streamlit`
+```
+pip install -r requirements.txt
+```
 
-`sdk_version` : _string_  
-Only applicable for `streamlit` SDK.  
-See [doc](https://hf.co/docs/hub/spaces) for more info on supported versions.
+Please refer to https://github.com/NVlabs/stylegan2-ada-pytorch for additional software/hardware requirements.
+
+## Dataset
+We follow the same dataset format as StyleGAN2-ADA supported, which can be either an image folder, or a zipped file.
+
+
+## Train a new StyleNeRF model
+```bash
+python run_train.py outdir=${OUTDIR} data=${DATASET} spec=paper512 model=stylenerf_ffhq
+```
+It will automatically detect all usable GPUs.
+
+Please check configuration files at ```conf/model``` and ```conf/spec```. You can always add your own model config. More details on how to use hydra configuration please follow https://hydra.cc/docs/intro/.
+
+## Render the pretrained model
+```bash
+python generate.py --outdir=${OUTDIR} --trunc=0.7 --seeds=${SEEDS} --network=${CHECKPOINT_PATH} --render-program="rotation_camera"
+```
+It supports different rotation trajectories for rendering new videos.
+
+## Run a demo page
+```bash
+python web_demo.py 21111
+```
+It will in default run a Gradio-powered demo on https://localhost:21111
+![Web demo](./docs/web_demo.gif)
+## Run a GUI visualizer
+```bash
+python visualizer.py
+```
+An interative application will show up for users to play with.
+![GUI demo](./docs/gui_demo.gif)
+## Citation
+
+```
+@inproceedings{
+    gu2022stylenerf,
+    title={StyleNeRF: A Style-based 3D Aware Generator for High-resolution Image Synthesis},
+    author={Jiatao Gu and Lingjie Liu and Peng Wang and Christian Theobalt},
+    booktitle={International Conference on Learning Representations},
+    year={2022},
+    url={https://openreview.net/forum?id=iUuzzTMUw9K}
+}
+```
+
+
+## License
+
+Copyright &copy; Facebook, Inc. All Rights Reserved.
+
+The majority of StyleNeRF is licensed under [CC-BY-NC](https://creativecommons.org/licenses/by-nc/4.0/), however, portions of this project are available under a separate license terms: all codes used or modified from [stylegan2-ada-pytorch](https://github.com/NVlabs/stylegan2-ada-pytorch) is under the [Nvidia Source Code License](https://nvlabs.github.io/stylegan2-ada-pytorch/license.html).
 
-`app_file`: _string_  
-Path to your main application file (which contains either `gradio` or `streamlit` Python code).  
-Path is relative to the root of the repository.
 
-`pinned`: _boolean_  
-Whether the Space stays on top of your list.

app.py

@@ -0,0 +1,213 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+import os, sys
+os.system('pip install -r requirements.txt')
+
+import gradio as gr
+import numpy as np
+import dnnlib
+import time
+import legacy
+import torch
+import glob
+
+import cv2
+import signal
+from torch_utils import misc
+from renderer import Renderer
+from training.networks import Generator
+from huggingface_hub import hf_hub_download
+
+
+device = torch.device('cuda')
+port   = int(sys.argv[1]) if len(sys.argv) > 1 else 21111
+
+
+
+def handler(signum, frame):
+    res = input("Ctrl-c was pressed. Do you really want to exit? y/n ")
+    if res == 'y':
+        gr.close_all()
+        exit(1)
+ 
+signal.signal(signal.SIGINT, handler)
+
+
+def set_random_seed(seed):
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+
+
+def get_camera_traj(model, pitch, yaw, fov=12, batch_size=1, model_name='FFHQ512'):
+    gen = model.synthesis
+    range_u, range_v = gen.C.range_u, gen.C.range_v
+    if not (('car' in model_name) or ('Car' in model_name)):  # TODO: hack, better option?
+        yaw, pitch = 0.5 * yaw, 0.3  * pitch
+        pitch = pitch + np.pi/2
+        u = (yaw - range_u[0]) / (range_u[1] - range_u[0])
+        v = (pitch - range_v[0]) / (range_v[1] - range_v[0])
+    else:
+        u = (yaw + 1) / 2
+        v = (pitch + 1) / 2
+    cam = gen.get_camera(batch_size=batch_size, mode=[u, v, 0.5], device=device, fov=fov)
+    return cam
+
+
+def check_name(model_name='FFHQ512'):
+    """Gets model by name."""
+    if model_name == 'FFHQ512':
+        network_pkl = hf_hub_download(repo_id='thomagram/stylenerf-ffhq-config-basic', filename='ffhq_512.pkl')
+    
+    # TODO: checkpoint to be updated!
+    # elif model_name == 'FFHQ512v2':
+    #     network_pkl = "./pretrained/ffhq_512_eg3d.pkl"
+    # elif model_name == 'AFHQ512':
+    #     network_pkl = "./pretrained/afhq_512.pkl"
+    # elif model_name == 'MetFaces512':
+    #     network_pkl = "./pretrained/metfaces_512.pkl"
+    # elif model_name == 'CompCars256':
+    #     network_pkl = "./pretrained/cars_256.pkl"
+    # elif model_name == 'FFHQ1024':
+    #     network_pkl = "./pretrained/ffhq_1024.pkl"
+    else:
+        if os.path.isdir(model_name):
+            network_pkl = sorted(glob.glob(model_name + '/*.pkl'))[-1]
+        else:
+            network_pkl = model_name
+    return network_pkl
+
+
+def get_model(network_pkl, render_option=None):
+    print('Loading networks from "%s"...' % network_pkl)
+    with dnnlib.util.open_url(network_pkl) as f:
+        network = legacy.load_network_pkl(f)
+        G = network['G_ema'].to(device)  # type: ignore
+
+    with torch.no_grad():
+        G2 = Generator(*G.init_args, **G.init_kwargs).to(device)
+        misc.copy_params_and_buffers(G, G2, require_all=False)
+
+    print('compile and go through the initial image')
+    G2 = G2.eval()
+    init_z = torch.from_numpy(np.random.RandomState(0).rand(1, G2.z_dim)).to(device)
+    init_cam = get_camera_traj(G2, 0, 0, model_name=network_pkl)
+    dummy = G2(z=init_z, c=None, camera_matrices=init_cam, render_option=render_option, theta=0)
+    res = dummy['img'].shape[-1]
+    imgs = np.zeros((res, res//2, 3))
+    return G2, res, imgs
+
+
+global_states = list(get_model(check_name()))
+wss  = [None, None]
+
+def proc_seed(history, seed):
+    if isinstance(seed, str):
+        seed = 0
+    else:
+        seed = int(seed)
+
+
+def f_synthesis(model_name, model_find, render_option, early, trunc, seed1, seed2, mix1, mix2, yaw, pitch, roll, fov, history):
+    history = history or {}
+    seeds = []
+    
+    if model_find != "":
+        model_name = model_find
+        
+    model_name = check_name(model_name)
+    if model_name != history.get("model_name", None):
+        model, res, imgs = get_model(model_name, render_option)
+        global_states[0] = model
+        global_states[1] = res
+        global_states[2] = imgs
+
+    model, res, imgs = global_states
+    for idx, seed in enumerate([seed1, seed2]):
+        if isinstance(seed, str):
+            seed = 0
+        else:
+            seed = int(seed)
+
+        if (seed != history.get(f'seed{idx}', -1)) or \
+            (model_name != history.get("model_name", None)) or \
+            (trunc != history.get("trunc", 0.7)) or \
+            (wss[idx] is None):
+            print(f'use seed {seed}')
+            set_random_seed(seed)
+            z   = torch.from_numpy(np.random.RandomState(int(seed)).randn(1, model.z_dim).astype('float32')).to(device)
+            ws  = model.mapping(z=z, c=None, truncation_psi=trunc)
+            img = model.get_final_output(styles=ws, camera_matrices=get_camera_traj(model, 0, 0), render_option=render_option)
+            ws  = ws.detach().cpu().numpy()
+            img = img[0].permute(1,2,0).detach().cpu().numpy()
+
+            
+            imgs[idx * res // 2: (1 + idx) * res // 2] = cv2.resize(
+                np.asarray(img).clip(-1, 1) * 0.5 + 0.5,
+                (res//2, res//2), cv2.INTER_AREA)
+            wss[idx] = ws
+        else:
+            seed = history[f'seed{idx}']
+        seeds += [seed]
+
+        history[f'seed{idx}'] = seed
+    history['trunc'] = trunc
+    history['model_name'] = model_name
+    
+    set_random_seed(sum(seeds))
+
+    # style mixing (?)
+    ws1, ws2 = [torch.from_numpy(ws).to(device) for ws in wss]
+    ws = ws1.clone()
+    ws[:, :8] = ws1[:, :8] * mix1 + ws2[:, :8] * (1 - mix1)
+    ws[:, 8:] = ws1[:, 8:] * mix2 + ws2[:, 8:] * (1 - mix2)
+    
+    # set visualization for other types of inputs.
+    if early == 'Normal Map':
+        render_option += ',normal,early'
+    elif early == 'Gradient Map':
+        render_option += ',gradient,early'
+    
+    start_t = time.time()
+    with torch.no_grad():
+        cam = get_camera_traj(model, pitch, yaw, fov, model_name=model_name)
+        image = model.get_final_output(
+            styles=ws, camera_matrices=cam, 
+            theta=roll * np.pi,
+            render_option=render_option)
+    end_t = time.time()
+
+    image = image[0].permute(1,2,0).detach().cpu().numpy().clip(-1, 1) * 0.5 + 0.5
+
+    if imgs.shape[0] == image.shape[0]:
+        image = np.concatenate([imgs, image], 1)
+    else:
+        a = image.shape[0]
+        b = int(imgs.shape[1] / imgs.shape[0] * a)
+        print(f'resize {a} {b} {image.shape} {imgs.shape}')
+        image = np.concatenate([cv2.resize(imgs, (b, a), cv2.INTER_AREA), image], 1)
+  
+    print(f'rendering time = {end_t-start_t:.4f}s')
+    image = (image * 255).astype('uint8')
+    return image, history
+
+model_name = gr.inputs.Dropdown(['FFHQ512'])  #  'FFHQ512v2', 'AFHQ512', 'MetFaces512', 'CompCars256', 'FFHQ1024'
+model_find = gr.inputs.Textbox(label="checkpoint path", default="")
+render_option = gr.inputs.Textbox(label="rendering options", default='steps:40')
+trunc  = gr.inputs.Slider(default=0.7, maximum=1.0, minimum=0.0, label='truncation trick')
+seed1  = gr.inputs.Number(default=1, label="seed1")
+seed2  = gr.inputs.Number(default=9, label="seed2")
+mix1   = gr.inputs.Slider(minimum=0, maximum=1, default=0, label="linear mixing ratio (geometry)")
+mix2   = gr.inputs.Slider(minimum=0, maximum=1, default=0, label="linear mixing ratio (apparence)")
+early  = gr.inputs.Radio(['None', 'Normal Map', 'Gradient Map'], default='None', label='intermedia output')
+yaw    = gr.inputs.Slider(minimum=-1, maximum=1, default=0, label="yaw")
+pitch  = gr.inputs.Slider(minimum=-1, maximum=1, default=0, label="pitch")
+roll   = gr.inputs.Slider(minimum=-1, maximum=1, default=0, label="roll (optional, not suggested)")
+fov    = gr.inputs.Slider(minimum=9, maximum=15, default=12, label="fov")
+css = ".output_image {height: 40rem !important; width: 100% !important;}"
+
+gr.Interface(fn=f_synthesis,
+             inputs=[model_name, model_find, render_option, early, trunc, seed1, seed2, mix1, mix2, yaw, pitch, roll, fov, "state"],
+             title="Interctive Web Demo for StyleNeRF (ICLR 2022)",
+             outputs=["image", "state"],
+             layout='unaligned',
+             css=css, theme='dark-huggingface',
+             live=True).launch(server_port=port)

conf/config.yaml

@@ -0,0 +1,49 @@
+defaults:
+  - _self_
+  - model: default
+  - spec: paper512
+
+# general options
+outdir: ~
+dry_run: False
+debug: False
+resume_run: ~
+
+snap: 50    # Snapshot interval [default: 50 ticks]
+imgsnap: 10
+metrics: [ "fid50k_full" ]
+seed: 2
+num_fp16_res: 4
+auto: False
+
+# dataset
+data: ~
+resolution: ~
+cond: False
+subset: ~   # Train with only N images: <int>, default = all
+mirror: False
+
+# discriminator augmentation
+aug: noaug
+p: ~
+target: ~
+augpipe: ~
+
+# transfer learning
+resume: ~
+freezed: ~
+
+# performance options
+fp32: False
+nhwc: False
+allow_tf32: False
+nobench: False
+workers: 3
+
+launcher: "spawn"
+partition: ~
+comment: ~
+gpus: ~     # Number of GPUs to use [default: 1]
+port: ~
+nodes: ~
+timeout: ~

conf/hydra/local.yaml

@@ -0,0 +1,16 @@
+sweep:
+  dir: /checkpoint/${env:USER}/space/gan/${env:PREFIX}/${hydra.job.name}
+  subdir: ${hydra.job.num}
+launcher:
+  submitit_folder: ${hydra.sweep.dir}
+  timeout_min: 4320
+  cpus_per_task: 64
+  gpus_per_node: 8
+  tasks_per_node: 1
+  mem_gb: 400
+  nodes: 1
+  name: ${env:PREFIX}_${hydra.job.config_name}
+  # partition: devlab,learnlab,learnfair,scavenge
+  # constraint: volta32gb
+  # max_num_timeout: 30
+  # exclude: learnfair1381,learnfair5192,learnfair2304

conf/model/default.yaml

@@ -0,0 +1,35 @@
+# @package _group_
+name: default
+
+G_kwargs:
+    class_name: "training.networks.Generator"
+    z_dim: 512
+    w_dim: 512
+
+    mapping_kwargs:
+        num_layers: ${spec.map}
+
+    synthesis_kwargs:
+        num_fp16_res: ${num_fp16_res}
+        channel_base: ${spec.fmaps}
+        channel_max: 512
+        conv_clamp: 256
+        architecture: skip
+
+D_kwargs:
+    class_name: "training.networks.Discriminator"
+    epilogue_kwargs:
+        mbstd_group_size: ${spec.mbstd}
+
+    num_fp16_res: ${num_fp16_res}
+    channel_base: ${spec.fmaps}
+    channel_max: 512
+    conv_clamp: 256
+    architecture: resnet
+
+# loss kwargs
+loss_kwargs:
+    pl_batch_shrink: 2
+    pl_decay: 0.01
+    pl_weight: 2
+    style_mixing_prob: 0.9

conf/model/stylenerf_afhq.yaml

@@ -0,0 +1,108 @@
+# @package _group_
+name: stylenerf_afhq
+
+G_kwargs:
+    class_name: "training.networks.Generator"
+    z_dim: 512
+    w_dim: 512
+
+    mapping_kwargs:
+        num_layers: ${spec.map}
+
+    synthesis_kwargs:
+        # global settings
+        num_fp16_res: ${num_fp16_res}
+        channel_base: 1
+        channel_max: 1024
+        conv_clamp: 256
+        kernel_size: 1
+        architecture: skip
+        upsample_mode: "nn_cat"
+
+        z_dim_bg: 32
+        z_dim: 0
+        resolution_vol: 32
+        resolution_start: 32
+        rgb_out_dim: 256
+
+        use_noise: False
+        module_name: "training.stylenerf.NeRFSynthesisNetwork"
+        no_bbox: True
+        margin: 0
+        magnitude_ema_beta: 0.999
+
+        camera_kwargs:
+            range_v: [1.4157963267948965, 1.7257963267948966]
+            range_u: [-0.3, 0.3]
+            range_radius: [1.0, 1.0]
+            depth_range: [0.88, 1.12]
+            fov: 12
+            gaussian_camera: True
+            angular_camera: True
+            depth_transform:  ~
+            dists_normalized: False
+            ray_align_corner: False
+            bg_start: 0.5
+        
+        renderer_kwargs:
+            n_bg_samples: 4
+            n_ray_samples: 14
+            abs_sigma: False
+            hierarchical: True
+            no_background: False
+            
+        foreground_kwargs:
+            positional_encoding: "normal"
+            downscale_p_by: 1
+            use_style: "StyleGAN2"
+            predict_rgb: True
+            use_viewdirs: False
+
+        background_kwargs:
+            positional_encoding: "normal"
+            hidden_size: 64
+            n_blocks: 4
+            downscale_p_by: 1
+            skips: []
+            inverse_sphere: True
+            use_style: "StyleGAN2"
+            predict_rgb: True
+            use_viewdirs: False
+
+        upsampler_kwargs:
+            channel_base: ${model.G_kwargs.synthesis_kwargs.channel_base}
+            channel_max:  ${model.G_kwargs.synthesis_kwargs.channel_max}
+            no_2d_renderer: False
+            no_residual_img: False
+            block_reses: ~
+            shared_rgb_style: False
+            upsample_type: "bilinear"
+        
+        progressive: True
+
+        # reuglarization
+        n_reg_samples: 16
+        reg_full: True
+
+D_kwargs:
+    class_name: "training.stylenerf.Discriminator"
+    epilogue_kwargs:
+        mbstd_group_size: ${spec.mbstd}
+
+    num_fp16_res: ${num_fp16_res}
+    channel_base: ${spec.fmaps}
+    channel_max: 512
+    conv_clamp: 256
+    architecture: skip
+    progressive: ${model.G_kwargs.synthesis_kwargs.progressive}
+    lowres_head: ${model.G_kwargs.synthesis_kwargs.resolution_start}
+    upsample_type: "bilinear"
+    resize_real_early: True
+
+# loss kwargs
+loss_kwargs:
+    pl_batch_shrink: 2
+    pl_decay: 0.01
+    pl_weight: 2
+    style_mixing_prob: 0.9
+    curriculum: [500,5000]

conf/model/stylenerf_cars.yaml

@@ -0,0 +1,108 @@
+# @package _group_
+name: stylenerf_ffhq
+
+G_kwargs:
+    class_name: "training.networks.Generator"
+    z_dim: 512
+    w_dim: 512
+
+    mapping_kwargs:
+        num_layers: ${spec.map}
+
+    synthesis_kwargs:
+        # global settings
+        num_fp16_res: ${num_fp16_res}
+        channel_base: 1
+        channel_max: 1024
+        conv_clamp: 256
+        kernel_size: 1
+        architecture: skip
+        upsample_mode: "pixelshuffle"
+
+        z_dim_bg: 32
+        z_dim: 0
+        resolution_vol: 32
+        resolution_start: 32
+        rgb_out_dim: 256
+
+        use_noise: False
+        module_name: "training.stylenerf.NeRFSynthesisNetwork"
+        no_bbox: True
+        margin: 0
+        magnitude_ema_beta: 0.999
+
+        camera_kwargs:
+            range_v: [1.4157963267948965, 1.7257963267948966]
+            range_u: [-3.141592653589793, 3.141592653589793]
+            range_radius: [1.0, 1.0]
+            depth_range: [0.8, 1.2]
+            fov: 16
+            gaussian_camera: False
+            angular_camera: True
+            depth_transform:  ~
+            dists_normalized: False
+            ray_align_corner: False
+            bg_start: 0.5
+        
+        renderer_kwargs:
+            n_bg_samples: 4
+            n_ray_samples: 16
+            abs_sigma: False
+            hierarchical: True
+            no_background: False
+            
+        foreground_kwargs:
+            positional_encoding: "normal"
+            downscale_p_by: 1
+            use_style: "StyleGAN2"
+            predict_rgb: True
+            use_viewdirs: False
+
+        background_kwargs:
+            positional_encoding: "normal"
+            hidden_size: 64
+            n_blocks: 4
+            downscale_p_by: 1
+            skips: []
+            inverse_sphere: True
+            use_style: "StyleGAN2"
+            predict_rgb: True
+            use_viewdirs: False
+
+        upsampler_kwargs:
+            channel_base: ${model.G_kwargs.synthesis_kwargs.channel_base}
+            channel_max:  ${model.G_kwargs.synthesis_kwargs.channel_max}
+            no_2d_renderer: False
+            no_residual_img: False
+            block_reses: ~
+            shared_rgb_style: False
+            upsample_type: "bilinear"
+        
+        progressive: True
+
+        # reuglarization
+        n_reg_samples: 0
+        reg_full: False
+
+D_kwargs:
+    class_name: "training.stylenerf.Discriminator"
+    epilogue_kwargs:
+        mbstd_group_size: ${spec.mbstd}
+
+    num_fp16_res: ${num_fp16_res}
+    channel_base: ${spec.fmaps}
+    channel_max: 512
+    conv_clamp: 256
+    architecture: skip
+    progressive: ${model.G_kwargs.synthesis_kwargs.progressive}
+    lowres_head: ${model.G_kwargs.synthesis_kwargs.resolution_start}
+    upsample_type: "bilinear"
+    resize_real_early: True
+
+# loss kwargs
+loss_kwargs:
+    pl_batch_shrink: 2
+    pl_decay: 0.01
+    pl_weight: 2
+    style_mixing_prob: 0.9
+    curriculum: [500,5000]

conf/model/stylenerf_cars_debug.yaml

@@ -0,0 +1,105 @@
+# @package _group_
+name: stylenerf_ffhq
+
+G_kwargs:
+    class_name: "training.networks.Generator"
+    z_dim: 512
+    w_dim: 512
+
+    mapping_kwargs:
+        num_layers: ${spec.map}
+
+    synthesis_kwargs:
+        # global settings
+        num_fp16_res: ${num_fp16_res}
+        channel_base: 1
+        channel_max: 1024
+        conv_clamp: 256
+        kernel_size: 1
+        architecture: skip
+        upsample_mode: "pixelshuffle"
+
+        z_dim_bg: 0
+        z_dim: 0
+        resolution_vol: 128
+        resolution_start: 32
+        rgb_out_dim: 256
+
+        use_noise: False
+        module_name: "training.stylenerf.NeRFSynthesisNetwork"
+        no_bbox: True
+        margin: 0
+        magnitude_ema_beta: 0.999
+
+        camera_kwargs:
+            range_v: [1.4157963267948965, 1.7257963267948966]
+            range_u: [-3.141592653589793, 3.141592653589793]
+            range_radius: [1.0, 1.0]
+            depth_range: [0.8, 1.2]
+            fov: 16
+            gaussian_camera: False
+            angular_camera: True
+            depth_transform:  ~
+            dists_normalized: False
+            ray_align_corner: False
+            bg_start: 0.5
+        
+        renderer_kwargs:
+            n_bg_samples: 0
+            n_ray_samples: 32
+            abs_sigma: False
+            hierarchical: True
+            no_background: True
+            
+        foreground_kwargs:
+            downscale_p_by: 1
+            use_style: "StyleGAN2"
+            predict_rgb: False
+            add_rgb: True
+            use_viewdirs: False
+            n_blocks: 0
+
+        input_kwargs:
+            output_mode: 'tri_plane_reshape'
+            input_mode: 'random'
+            in_res:  4
+            out_res: 256
+            out_dim: 32
+
+        upsampler_kwargs:
+            channel_base: ${model.G_kwargs.synthesis_kwargs.channel_base}
+            channel_max:  ${model.G_kwargs.synthesis_kwargs.channel_max}
+            no_2d_renderer: False
+            no_residual_img: False
+            block_reses: ~
+            shared_rgb_style: False
+            upsample_type: "bilinear"
+        
+        progressive: True
+
+        # reuglarization
+        n_reg_samples: 0
+        reg_full: False
+
+D_kwargs:
+    class_name: "training.stylenerf.Discriminator"
+    epilogue_kwargs:
+        mbstd_group_size: ${spec.mbstd}
+
+    num_fp16_res: ${num_fp16_res}
+    channel_base: ${spec.fmaps}
+    channel_max: 512
+    conv_clamp: 256
+    architecture: skip
+    progressive: ${model.G_kwargs.synthesis_kwargs.progressive}
+    lowres_head: ${model.G_kwargs.synthesis_kwargs.resolution_start}
+    upsample_type: "bilinear"
+    resize_real_early: True
+
+# loss kwargs
+loss_kwargs:
+    pl_batch_shrink: 2
+    pl_decay: 0.01
+    pl_weight: 2
+    style_mixing_prob: 0.9
+    curriculum: [500,5000]

conf/model/stylenerf_ffhq.yaml

@@ -0,0 +1,108 @@
+# @package _group_
+name: stylenerf_ffhq
+
+G_kwargs:
+    class_name: "training.networks.Generator"
+    z_dim: 512
+    w_dim: 512
+
+    mapping_kwargs:
+        num_layers: ${spec.map}
+
+    synthesis_kwargs:
+        # global settings
+        num_fp16_res: ${num_fp16_res}
+        channel_base: 1
+        channel_max: 1024
+        conv_clamp: 256
+        kernel_size: 1
+        architecture: skip
+        upsample_mode: "nn_cat"
+
+        z_dim_bg: 32
+        z_dim: 0
+        resolution_vol: 32
+        resolution_start: 32
+        rgb_out_dim: 256
+
+        use_noise: False
+        module_name: "training.stylenerf.NeRFSynthesisNetwork"
+        no_bbox: True
+        margin: 0
+        magnitude_ema_beta: 0.999
+
+        camera_kwargs:
+            range_v: [1.4157963267948965, 1.7257963267948966]
+            range_u: [-0.3, 0.3]
+            range_radius: [1.0, 1.0]
+            depth_range: [0.88, 1.12]
+            fov: 12
+            gaussian_camera: True
+            angular_camera: True
+            depth_transform:  ~
+            dists_normalized: False
+            ray_align_corner: False
+            bg_start: 0.5
+        
+        renderer_kwargs:
+            n_bg_samples: 4
+            n_ray_samples: 14
+            abs_sigma: False
+            hierarchical: True
+            no_background: False
+            
+        foreground_kwargs:
+            positional_encoding: "normal"
+            downscale_p_by: 1
+            use_style: "StyleGAN2"
+            predict_rgb: True
+            use_viewdirs: False
+
+        background_kwargs:
+            positional_encoding: "normal"
+            hidden_size: 64
+            n_blocks: 4
+            downscale_p_by: 1
+            skips: []
+            inverse_sphere: True
+            use_style: "StyleGAN2"
+            predict_rgb: True
+            use_viewdirs: False
+
+        upsampler_kwargs:
+            channel_base: ${model.G_kwargs.synthesis_kwargs.channel_base}
+            channel_max:  ${model.G_kwargs.synthesis_kwargs.channel_max}
+            no_2d_renderer: False
+            no_residual_img: False
+            block_reses: ~
+            shared_rgb_style: False
+            upsample_type: "bilinear"
+        
+        progressive: True
+
+        # reuglarization
+        n_reg_samples: 16
+        reg_full: True
+
+D_kwargs:
+    class_name: "training.stylenerf.Discriminator"
+    epilogue_kwargs:
+        mbstd_group_size: ${spec.mbstd}
+
+    num_fp16_res: ${num_fp16_res}
+    channel_base: ${spec.fmaps}
+    channel_max: 512
+    conv_clamp: 256
+    architecture: skip
+    progressive: ${model.G_kwargs.synthesis_kwargs.progressive}
+    lowres_head: ${model.G_kwargs.synthesis_kwargs.resolution_start}
+    upsample_type: "bilinear"
+    resize_real_early: True
+
+# loss kwargs
+loss_kwargs:
+    pl_batch_shrink: 2
+    pl_decay: 0.01
+    pl_weight: 2
+    style_mixing_prob: 0.9
+    curriculum: [500,5000]

conf/model/stylenerf_ffhq_ae.yaml

@@ -0,0 +1,118 @@
+# @package _group_
+name: stylenerf_ffhq
+
+G_kwargs:
+    class_name: "training.networks.Generator"
+    z_dim: 512
+    w_dim: 512
+
+    mapping_kwargs:
+        num_layers: ${spec.map}
+
+    synthesis_kwargs:
+        # global settings
+        num_fp16_res: ${num_fp16_res}
+        channel_base: 1
+        channel_max: 1024
+        conv_clamp: 256
+        kernel_size: 1
+        architecture: skip
+        upsample_mode: "nn_cat"
+
+        z_dim: 0
+        resolution_vol: 128
+        resolution_start: 128
+        rgb_out_dim: 32
+
+        use_noise: False
+        module_name: "training.stylenerf.NeRFSynthesisNetwork"
+        no_bbox: True
+        margin: 0
+        magnitude_ema_beta: 0.999
+
+        camera_kwargs:
+            range_v: [1.4157963267948965, 1.7257963267948966]
+            range_u: [-0.3, 0.3]
+            range_radius:  [1.0, 1.0]
+            depth_range: [0.88, 1.12]
+            fov: 12
+            gaussian_camera: True
+            angular_camera: True
+            depth_transform: ~
+            dists_normalized: True
+            ray_align_corner: False
+            bg_start: 0.5
+        
+        renderer_kwargs:
+            n_ray_samples: 32
+            abs_sigma: False
+            hierarchical: True
+            no_background: True
+            
+        foreground_kwargs:
+            downscale_p_by: 1
+            use_style: "StyleGAN2"
+            predict_rgb: False
+            use_viewdirs: False
+            add_rgb: True
+            n_blocks: 0
+
+        input_kwargs:
+            output_mode: 'tri_plane_reshape'
+            input_mode: 'random'
+            in_res:  4
+            out_res: 256
+            out_dim: 32
+
+        upsampler_kwargs:
+            no_2d_renderer: False
+            no_residual_img: False
+            block_reses: ~
+            shared_rgb_style: False
+            upsample_type: "bilinear"
+        
+        progressive: True
+
+        # reuglarization
+        n_reg_samples: 0
+        reg_full: False
+
+    encoder_kwargs:
+        class_name: "training.stylenerf.Encoder"
+        num_fp16_res: ${num_fp16_res}
+        channel_base: ${spec.fmaps}
+        channel_max: 512
+        conv_clamp: 256
+        architecture: skip
+        progressive: ${..synthesis_kwargs.progressive}
+        lowres_head: ${..synthesis_kwargs.resolution_start}
+        upsample_type: "bilinear"
+        model_kwargs:
+            output_mode: "W+"
+            predict_camera: False
+
+D_kwargs:
+    class_name: "training.stylenerf.Discriminator"
+    epilogue_kwargs:
+        mbstd_group_size: ${spec.mbstd}
+
+    num_fp16_res: ${num_fp16_res}
+    channel_base: ${spec.fmaps}
+    channel_max: 512
+    conv_clamp: 256
+    architecture: skip
+    
+    predict_camera: True
+
+    progressive: ${model.G_kwargs.synthesis_kwargs.progressive}
+    lowres_head: ${model.G_kwargs.synthesis_kwargs.resolution_start}
+    upsample_type: "bilinear"
+    resize_real_early: True
+
+# loss kwargs
+loss_kwargs:
+    pl_batch_shrink: 2
+    pl_decay: 0.01
+    pl_weight: 2
+    style_mixing_prob: 0.9
+    curriculum: [500,5000]

conf/model/stylenerf_ffhq_ae_basic.yaml

@@ -0,0 +1,110 @@
+# @package _group_
+name: stylenerf_ffhq
+
+G_kwargs:
+    class_name: "training.networks.Generator"
+    z_dim: 512
+    w_dim: 512
+
+    mapping_kwargs:
+        num_layers: ${spec.map}
+
+    synthesis_kwargs:
+        # global settings
+        num_fp16_res: ${num_fp16_res}
+        channel_base: 1
+        channel_max: 1024
+        conv_clamp: 256
+        kernel_size: 1
+        architecture: skip
+        upsample_mode: "nn_cat"
+
+        z_dim: 0
+        resolution_vol: 32
+        resolution_start: 32
+        rgb_out_dim: 32
+
+        use_noise: False
+        module_name: "training.stylenerf.NeRFSynthesisNetwork"
+        no_bbox: True
+        margin: 0
+        magnitude_ema_beta: 0.999
+
+        camera_kwargs:
+            range_v: [1.4157963267948965, 1.7257963267948966]
+            range_u: [-0.3, 0.3]
+            range_radius:  [1.0, 1.0]
+            depth_range: [0.88, 1.12]
+            fov: 12
+            gaussian_camera: True
+            angular_camera: True
+            depth_transform: ~
+            dists_normalized: True
+            ray_align_corner: False
+            bg_start: 0.5
+        
+        renderer_kwargs:
+            n_ray_samples: 32
+            abs_sigma: False
+            hierarchical: True
+            no_background: True
+            
+        foreground_kwargs:
+            downscale_p_by: 1
+            use_style: "StyleGAN2"
+            predict_rgb: False
+            use_viewdirs: False
+            add_rgb: True
+
+        upsampler_kwargs:
+            no_2d_renderer: False
+            no_residual_img: False
+            block_reses: ~
+            shared_rgb_style: False
+            upsample_type: "bilinear"
+        
+        progressive: True
+
+        # reuglarization
+        n_reg_samples: 0
+        reg_full: False
+
+    encoder_kwargs:
+        class_name: "training.stylenerf.Encoder"
+        num_fp16_res: ${num_fp16_res}
+        channel_base: ${spec.fmaps}
+        channel_max: 512
+        conv_clamp: 256
+        architecture: skip
+        progressive: ${..synthesis_kwargs.progressive}
+        lowres_head: ${..synthesis_kwargs.resolution_start}
+        upsample_type: "bilinear"
+        model_kwargs:
+            output_mode: "W+"
+            predict_camera: False
+
+D_kwargs:
+    class_name: "training.stylenerf.Discriminator"
+    epilogue_kwargs:
+        mbstd_group_size: ${spec.mbstd}
+
+    num_fp16_res: ${num_fp16_res}
+    channel_base: ${spec.fmaps}
+    channel_max: 512
+    conv_clamp: 256
+    architecture: skip
+    
+    predict_camera: True
+
+    progressive: ${model.G_kwargs.synthesis_kwargs.progressive}
+    lowres_head: ${model.G_kwargs.synthesis_kwargs.resolution_start}
+    upsample_type: "bilinear"
+    resize_real_early: True
+
+# loss kwargs
+loss_kwargs:
+    pl_batch_shrink: 2
+    pl_decay: 0.01
+    pl_weight: 2
+    style_mixing_prob: 0.9
+    curriculum: [500,5000]

conf/model/stylenerf_ffhq_debug.yaml

@@ -0,0 +1,103 @@
+# @package _group_
+name: stylenerf_ffhq
+
+G_kwargs:
+    class_name: "training.networks.Generator"
+    z_dim: 512
+    w_dim: 512
+
+    mapping_kwargs:
+        num_layers: ${spec.map}
+
+    synthesis_kwargs:
+        # global settings
+        num_fp16_res: ${num_fp16_res}
+        channel_base: 1
+        channel_max: 1024
+        conv_clamp: 256
+        kernel_size: 1
+        architecture: skip
+        upsample_mode: "nn_cat"
+
+        z_dim: 0
+        resolution_vol: 128
+        resolution_start: 128
+        rgb_out_dim: 32
+
+        use_noise: False
+        module_name: "training.stylenerf.NeRFSynthesisNetwork"
+        no_bbox: True
+        margin: 0
+        magnitude_ema_beta: 0.999
+
+        camera_kwargs:
+            range_v: [1.4157963267948965, 1.7257963267948966]
+            range_u: [-0.3, 0.3]
+            range_radius:  [1.0, 1.0]
+            depth_range: [0.88, 1.12]
+            fov: 12
+            gaussian_camera: True
+            angular_camera: True
+            depth_transform: ~
+            dists_normalized: True
+            ray_align_corner: False
+            bg_start: 0.5
+        
+        renderer_kwargs:
+            n_ray_samples: 32
+            abs_sigma: False
+            hierarchical: True
+            no_background: True
+            
+        foreground_kwargs:
+            downscale_p_by: 1
+            use_style: "StyleGAN2"
+            predict_rgb: False
+            use_viewdirs: False
+            add_rgb: True
+            n_blocks: 0
+
+        input_kwargs:
+            output_mode: 'tri_plane_reshape'
+            input_mode: 'random'
+            in_res:  4
+            out_res: 256
+            out_dim: 32
+            keep_posenc: -1
+            keep_nerf_latents: False
+
+        upsampler_kwargs:
+            no_2d_renderer: False
+            no_residual_img: False
+            block_reses: ~
+            shared_rgb_style: False
+            upsample_type: "bilinear"
+        
+        progressive: True
+
+        # reuglarization
+        n_reg_samples: 0
+        reg_full: False
+
+D_kwargs:
+    class_name: "training.stylenerf.Discriminator"
+    epilogue_kwargs:
+        mbstd_group_size: ${spec.mbstd}
+
+    num_fp16_res: ${num_fp16_res}
+    channel_base: ${spec.fmaps}
+    channel_max: 512
+    conv_clamp: 256
+    architecture: skip
+    progressive: ${model.G_kwargs.synthesis_kwargs.progressive}
+    lowres_head: ${model.G_kwargs.synthesis_kwargs.resolution_start}
+    upsample_type: "bilinear"
+    resize_real_early: True
+
+# loss kwargs
+loss_kwargs:
+    pl_batch_shrink: 2
+    pl_decay: 0.01
+    pl_weight: 2
+    style_mixing_prob: 0.9
+    curriculum: [500,5000]

conf/model/stylenerf_ffhq_eg3d.yaml

@@ -0,0 +1,100 @@
+# @package _group_
+name: stylenerf_ffhq
+
+G_kwargs:
+    class_name: "training.networks.Generator"
+    z_dim: 512
+    w_dim: 512
+
+    mapping_kwargs:
+        num_layers: ${spec.map}
+
+    synthesis_kwargs:
+        # global settings
+        num_fp16_res: ${num_fp16_res}
+        channel_base: 1
+        channel_max: 512
+        conv_clamp: 256
+        kernel_size: 3
+        architecture: skip
+
+        z_dim: 0
+        resolution_vol: 128
+        resolution_start: 32
+        rgb_out_dim: 32
+
+        use_noise: False
+        module_name: "training.stylenerf.NeRFSynthesisNetwork"
+        no_bbox: True
+        margin: 0
+        magnitude_ema_beta: 0.999
+
+        camera_kwargs:
+            range_v: [1.4157963267948965, 1.7257963267948966]
+            range_u: [-0.3, 0.3]
+            range_radius:  [1.0, 1.0]
+            depth_range: [0.88, 1.12]
+            fov: 12
+            gaussian_camera: True
+            angular_camera: True
+            depth_transform: ~
+            dists_normalized: True
+            ray_align_corner: False
+            bg_start: 0.5
+        
+        renderer_kwargs:
+            n_ray_samples: 32
+            abs_sigma: False
+            hierarchical: True
+            no_background: True
+            
+        foreground_kwargs:
+            downscale_p_by: 1
+            use_style: "StyleGAN2"
+            predict_rgb: False
+            use_viewdirs: False
+            add_rgb: True
+            n_blocks: 0
+
+        input_kwargs:
+            output_mode: 'tri_plane_reshape'
+            input_mode: 'random'
+            in_res:  4
+            out_res: 256
+            out_dim: 32
+
+        upsampler_kwargs:
+            no_2d_renderer: False
+            block_reses: ~
+            shared_rgb_style: False
+            upsample_type: "bilinear"
+        
+        progressive: False
+        prog_nerf_only: True
+
+        # reuglarization
+        n_reg_samples: 0
+        reg_full: False
+
+D_kwargs:
+    class_name: "training.stylenerf.Discriminator"
+    epilogue_kwargs:
+        mbstd_group_size: ${spec.mbstd}
+
+    num_fp16_res: ${num_fp16_res}
+    channel_base: ${spec.fmaps}
+    channel_max: 512
+    conv_clamp: 256
+    architecture: skip
+    progressive: False
+    dual_input_res: 128
+    upsample_type: "bilinear"
+    resize_real_early: True
+
+# loss kwargs
+loss_kwargs:
+    pl_batch_shrink: 2
+    pl_decay: 0.01
+    pl_weight: 2
+    style_mixing_prob: 0.9
+    curriculum: [0,5000]

conf/model/stylenerf_ffhq_warped_depth.yaml

@@ -0,0 +1,97 @@
+# @package _group_
+name: stylenerf_ffhq_warped_depth
+
+G_kwargs:
+    class_name: "training.networks.Generator"
+    z_dim: 512
+    w_dim: 512
+
+    mapping_kwargs:
+        num_layers: ${spec.map}
+
+    synthesis_kwargs:
+        # global settings
+        num_fp16_res: ${num_fp16_res}
+        channel_base: 1
+        channel_max: 1024
+        conv_clamp: 256
+        kernel_size: 1
+        architecture: skip
+        upsample_mode: "nn_cat"
+
+        z_dim_bg: 0
+        z_dim: 0
+        resolution_vol: 32
+        resolution_start: 32
+        rgb_out_dim: 256
+
+        use_noise: False
+        module_name: "training.stylenerf.NeRFSynthesisNetwork"
+        no_bbox: True
+        margin: 0
+        magnitude_ema_beta: 0.999
+
+        camera_kwargs:
+            range_v: [1.4157963267948965, 1.7257963267948966]
+            range_u: [-0.3, 0.3]
+            range_radius: [1.0, 1.0]
+            depth_range: [0.88, 3.2]
+            fov: 12
+            gaussian_camera: True
+            angular_camera: True
+            depth_transform:  InverseWarp
+            dists_normalized: True
+            ray_align_corner: False
+            bg_start: 0.5
+        
+        renderer_kwargs:
+            n_bg_samples: 0
+            n_ray_samples: 48
+            abs_sigma: False
+            hierarchical: True
+            no_background: True
+            
+        foreground_kwargs:
+            positional_encoding: "normal"
+            downscale_p_by: 1
+            use_style: "StyleGAN2"
+            predict_rgb: True
+            use_viewdirs: False
+
+        upsampler_kwargs:
+            channel_base: ${model.G_kwargs.synthesis_kwargs.channel_base}
+            channel_max:  ${model.G_kwargs.synthesis_kwargs.channel_max}
+            no_2d_renderer: False
+            no_residual_img: False
+            block_reses: ~
+            shared_rgb_style: False
+            upsample_type: "bilinear"
+        
+        progressive: True
+
+        # reuglarization
+        n_reg_samples: 16
+        reg_full: True
+
+D_kwargs:
+    class_name: "training.stylenerf.Discriminator"
+    epilogue_kwargs:
+        mbstd_group_size: ${spec.mbstd}
+
+    num_fp16_res: ${num_fp16_res}
+    channel_base: ${spec.fmaps}
+    channel_max: 512
+    conv_clamp: 256
+    architecture: skip
+    progressive: ${model.G_kwargs.synthesis_kwargs.progressive}
+    lowres_head: ${model.G_kwargs.synthesis_kwargs.resolution_start}
+    upsample_type: "bilinear"
+    resize_real_early: True
+
+# loss kwargs
+loss_kwargs:
+    pl_batch_shrink: 2
+    pl_decay: 0.01
+    pl_weight: 2
+    style_mixing_prob: 0.9
+    curriculum: [500,5000]

conf/spec/cifar.yaml

@@ -0,0 +1,13 @@
+# @package _group_
+
+name: cifar
+ref_gpus: 2
+kimg: 100000
+mb: 64
+mbstd: 32
+fmaps: 1
+lrate: 0.0025
+gamma: 0.01
+ema: 500
+ramp: 0.05
+map: 2

conf/spec/nerf32.yaml

@@ -0,0 +1,14 @@
+# @package _group_
+
+name: nerf32
+ref_gpus: 8
+kimg: 25000
+mb: 64
+mbstd: 8
+fmaps: 0.5
+lrate: 0.0025
+lrate_disc: 0.0025
+gamma: 0.003
+ema: 20
+ramp: ~
+map: 8

conf/spec/paper1024.yaml

@@ -0,0 +1,14 @@
+# @package _group_
+
+name: paper1024
+ref_gpus: 8
+kimg: 25000
+mb: 32
+mbstd: 4
+fmaps: 1
+lrate: 0.002
+lrate_disc: 0.002
+gamma: 2
+ema: 10
+ramp: ~
+map: 8

conf/spec/paper256.yaml

@@ -0,0 +1,14 @@
+# @package _group_
+
+name: paper256
+ref_gpus: 8
+kimg: 25000
+mb: 64
+mbstd: 8
+fmaps: 0.5
+lrate: 0.0025
+lrate_disc: 0.0025
+gamma: 0.5
+ema: 20
+ramp: ~
+map: 8

conf/spec/paper512.yaml

@@ -0,0 +1,12 @@
+name: paper512
+ref_gpus: 8
+kimg: 25000
+mb: 64
+mbstd: 8
+fmaps: 1
+lrate: 0.0025
+lrate_disc: 0.0025
+gamma: 0.5
+ema: 20
+ramp: ~
+map: 8

conf/spec/stylegan2.yaml

@@ -0,0 +1,14 @@
+# @package _group_
+
+name: stylegan2
+ref_gpus: 8
+kimg: 25000
+mb: 32
+mbstd: 4
+fmaps: 1
+lrate: 0.002
+lrate_disc: 0.0025
+gamma: 10
+ema: 10
+ramp: ~
+map: 8

dnnlib/__init__.py

@@ -0,0 +1,9 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+from .util import EasyDict, make_cache_dir_path

dnnlib/camera.py

@@ -0,0 +1,687 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+
+import numpy as np
+from numpy.lib.function_base import angle
+import torch
+import torch.nn.functional as F
+import math
+
+from scipy.spatial.transform import Rotation as Rot
+HUGE_NUMBER = 1e10
+TINY_NUMBER = 1e-6      # float32 only has 7 decimal digits precision
+
+
+def get_camera_mat(fov=49.13, invert=True):
+    # fov = 2 * arctan(sensor / (2 * focal))
+    # focal = (sensor / 2)  * 1 / (tan(0.5 * fov))
+    # in our case, sensor = 2 as pixels are in [-1, 1]
+    focal = 1. / np.tan(0.5 * fov * np.pi/180.)
+    focal = focal.astype(np.float32)
+    mat = torch.tensor([
+        [focal, 0., 0., 0.],
+        [0., focal, 0., 0.],
+        [0., 0., 1, 0.],
+        [0., 0., 0., 1.]
+    ]).reshape(1, 4, 4)
+    if invert:
+        mat = torch.inverse(mat)
+    return mat
+
+
+def get_random_pose(range_u, range_v, range_radius, batch_size=32,
+                    invert=False, gaussian=False, angular=False):
+    loc, (u, v) = sample_on_sphere(range_u, range_v, size=(batch_size), gaussian=gaussian, angular=angular)
+    radius = range_radius[0] + torch.rand(batch_size) * (range_radius[1] - range_radius[0])
+    loc = loc * radius.unsqueeze(-1)
+    R = look_at(loc)
+    RT = torch.eye(4).reshape(1, 4, 4).repeat(batch_size, 1, 1)
+    RT[:, :3, :3] = R
+    RT[:, :3, -1] = loc
+
+    if invert:
+        RT = torch.inverse(RT)
+    
+    def N(a, range_a):
+        if range_a[0] == range_a[1]:
+            return a * 0
+        return (a - range_a[0]) / (range_a[1] - range_a[0])
+    
+    val_u, val_v, val_r = N(u, range_u), N(v, range_v), N(radius, range_radius)
+    return RT, (val_u, val_v, val_r)
+
+
+def get_camera_pose(range_u, range_v, range_r, val_u=0.5, val_v=0.5, val_r=0.5,
+                    batch_size=32, invert=False,  gaussian=False, angular=False):
+    r0, rr = range_r[0], range_r[1] - range_r[0]
+    r = r0 + val_r * rr
+    if not gaussian:
+        u0, ur = range_u[0], range_u[1] - range_u[0]
+        v0, vr = range_v[0], range_v[1] - range_v[0]   
+        u = u0 + val_u * ur
+        v = v0 + val_v * vr
+    else:
+        mean_u, mean_v = sum(range_u) / 2, sum(range_v) / 2
+        vu, vv = mean_u - range_u[0], mean_v - range_v[0]
+        u = mean_u + vu * val_u
+        v = mean_v + vv * val_v
+        
+    loc, _ = sample_on_sphere((u, u), (v, v), size=(batch_size), angular=angular)
+    radius = torch.ones(batch_size) * r
+    loc = loc * radius.unsqueeze(-1)
+    R = look_at(loc)
+    RT = torch.eye(4).reshape(1, 4, 4).repeat(batch_size, 1, 1)
+    RT[:, :3, :3] = R
+    RT[:, :3, -1] = loc
+
+    if invert:
+        RT = torch.inverse(RT)
+    return RT
+
+
+def get_camera_pose_v2(range_u, range_v, range_r, mode, invert=False, gaussian=False, angular=False):
+    r0, rr = range_r[0], range_r[1] - range_r[0]
+    val_u, val_v = mode[:,0], mode[:,1]
+    val_r = torch.ones_like(val_u) * 0.5
+    if not gaussian:
+        u0, ur = range_u[0], range_u[1] - range_u[0]
+        v0, vr = range_v[0], range_v[1] - range_v[0]
+        u = u0 + val_u * ur
+        v = v0 + val_v * vr
+    else:
+        mean_u, mean_v = sum(range_u) / 2, sum(range_v) / 2
+        vu, vv = mean_u - range_u[0], mean_v - range_v[0]
+        u = mean_u + vu * val_u
+        v = mean_v + vv * val_v
+    
+    loc = to_sphere(u, v, angular)
+    radius = r0 + val_r * rr
+    loc = loc * radius.unsqueeze(-1)
+    R = look_at(loc)
+    RT = torch.eye(4).to(R.device).reshape(1, 4, 4).repeat(R.size(0), 1, 1)
+    RT[:, :3, :3] = R
+    RT[:, :3, -1] = loc
+
+    if invert:
+        RT = torch.inverse(RT)
+    return RT, (val_u, val_v, val_r)
+
+
+def to_sphere(u, v, angular=False):
+    T = torch if isinstance(u, torch.Tensor) else np
+    if not angular:
+        theta = 2 * math.pi * u
+        phi = T.arccos(1 - 2 * v)
+    else:
+        theta, phi = u, v
+    
+    cx = T.sin(phi) * T.cos(theta)
+    cy = T.sin(phi) * T.sin(theta)
+    cz = T.cos(phi)
+    return T.stack([cx, cy, cz], -1)
+
+
+def sample_on_sphere(range_u=(0, 1), range_v=(0, 1), size=(1,),
+                     to_pytorch=True, gaussian=False, angular=False):
+    if not gaussian:
+        u = np.random.uniform(*range_u, size=size)
+        v = np.random.uniform(*range_v, size=size)
+    else:
+        mean_u, mean_v = sum(range_u) / 2, sum(range_v) / 2
+        var_u, var_v = mean_u - range_u[0], mean_v - range_v[0]
+        u = np.random.normal(size=size) * var_u + mean_u
+        v = np.random.normal(size=size) * var_v + mean_v
+
+    sample = to_sphere(u, v, angular)
+    if to_pytorch:
+        sample = torch.tensor(sample).float()
+        u, v = torch.tensor(u).float(), torch.tensor(v).float()
+
+    return sample, (u, v)
+
+
+def look_at(eye, at=np.array([0, 0, 0]), up=np.array([0, 0, 1]), eps=1e-5,
+            to_pytorch=True):
+    if not isinstance(eye, torch.Tensor):
+        # this is the original code from GRAF
+        at = at.astype(float).reshape(1, 3)
+        up = up.astype(float).reshape(1, 3)
+        eye = eye.reshape(-1, 3)
+        up = up.repeat(eye.shape[0] // up.shape[0], axis=0)
+        eps = np.array([eps]).reshape(1, 1).repeat(up.shape[0], axis=0)
+        z_axis = eye - at
+        z_axis /= np.max(np.stack([np.linalg.norm(z_axis,
+                                                axis=1, keepdims=True), eps]))
+        x_axis = np.cross(up, z_axis)
+        x_axis /= np.max(np.stack([np.linalg.norm(x_axis,
+                                                axis=1, keepdims=True), eps]))
+        y_axis = np.cross(z_axis, x_axis)
+        y_axis /= np.max(np.stack([np.linalg.norm(y_axis,
+                                                axis=1, keepdims=True), eps]))
+        r_mat = np.concatenate(
+            (x_axis.reshape(-1, 3, 1), y_axis.reshape(-1, 3, 1), z_axis.reshape(
+                -1, 3, 1)), axis=2)
+        if to_pytorch:
+            r_mat = torch.tensor(r_mat).float()
+    else:
+        
+        def normalize(x, axis=-1, order=2):
+            l2 = x.norm(p=order, dim=axis, keepdim=True).clamp(min=1e-8)
+            return x / l2
+        
+        at, up = torch.from_numpy(at).float().to(eye.device), torch.from_numpy(up).float().to(eye.device)
+        z_axis = normalize(eye - at[None, :])
+        x_axis = normalize(torch.cross(up[None,:].expand_as(z_axis), z_axis, dim=-1))
+        y_axis = normalize(torch.cross(z_axis, x_axis, dim=-1))
+        r_mat = torch.stack([x_axis, y_axis, z_axis], dim=-1)
+
+    return r_mat
+
+
+def get_rotation_matrix(axis='z', value=0., batch_size=32):
+    r = Rot.from_euler(axis, value * 2 * np.pi).as_dcm()
+    r = torch.from_numpy(r).reshape(1, 3, 3).repeat(batch_size, 1, 1)
+    return r
+
+
+def get_corner_rays(corner_pixels, camera_matrices, res):
+    assert (res + 1) * (res + 1) == corner_pixels.size(1)
+    batch_size = camera_matrices[0].size(0)
+    rays, origins, _ = get_camera_rays(camera_matrices, corner_pixels)
+    corner_rays = torch.cat([rays, torch.cross(origins, rays, dim=-1)], -1)
+    corner_rays = corner_rays.reshape(batch_size, res+1, res+1, 6).permute(0,3,1,2)
+    corner_rays = torch.cat([corner_rays[..., :-1, :-1], corner_rays[..., 1:, :-1], corner_rays[..., 1:, 1:], corner_rays[..., :-1, 1:]], 1)
+    return corner_rays
+    
+
+def arange_pixels(
+        resolution=(128, 128), 
+        batch_size=1, 
+        subsample_to=None, 
+        invert_y_axis=False, 
+        margin=0,
+        corner_aligned=True,
+        jitter=None
+    ):
+    ''' Arranges pixels for given resolution in range image_range.
+
+    The function returns the unscaled pixel locations as integers and the
+    scaled float values.
+
+    Args:
+        resolution (tuple): image resolution
+        batch_size (int): batch size
+        subsample_to (int): if integer and > 0, the points are randomly
+            subsampled to this value
+    '''
+    h, w = resolution
+    n_points = resolution[0] * resolution[1]
+    uh = 1 if corner_aligned else 1 - (1 / h)
+    uw = 1 if corner_aligned else 1 - (1 / w)
+    if margin > 0:
+        uh = uh + (2 / h) * margin
+        uw = uw + (2 / w) * margin 
+        w, h = w + margin * 2, h + margin * 2
+
+    x, y = torch.linspace(-uw, uw, w), torch.linspace(-uh, uh, h)
+    if jitter is not None:
+        dx = (torch.ones_like(x).uniform_() - 0.5) * 2 / w * jitter
+        dy = (torch.ones_like(y).uniform_() - 0.5) * 2 / h * jitter
+        x, y = x + dx, y + dy
+    x, y = torch.meshgrid(x, y)
+    pixel_scaled = torch.stack([x, y], -1).permute(1,0,2).reshape(1, -1, 2).repeat(batch_size, 1, 1)
+    
+    # Subsample points if subsample_to is not None and > 0
+    if (subsample_to is not None and subsample_to > 0 and subsample_to < n_points):
+        idx = np.random.choice(pixel_scaled.shape[1], size=(subsample_to,),
+                               replace=False)
+        pixel_scaled = pixel_scaled[:, idx]
+
+    if invert_y_axis:
+        pixel_scaled[..., -1] *= -1.
+
+    return pixel_scaled
+
+
+def to_pytorch(tensor, return_type=False):
+    ''' Converts input tensor to pytorch.
+
+    Args:
+        tensor (tensor): Numpy or Pytorch tensor
+        return_type (bool): whether to return input type
+    '''
+    is_numpy = False
+    if type(tensor) == np.ndarray:
+        tensor = torch.from_numpy(tensor)
+        is_numpy = True
+    tensor = tensor.clone()
+    if return_type:
+        return tensor, is_numpy
+    return tensor
+
+
+def transform_to_world(pixels, depth, camera_mat, world_mat, scale_mat=None,
+                       invert=True, use_absolute_depth=True):
+    ''' Transforms pixel positions p with given depth value d to world coordinates.
+
+    Args:
+        pixels (tensor): pixel tensor of size B x N x 2
+        depth (tensor): depth tensor of size B x N x 1
+        camera_mat (tensor): camera matrix
+        world_mat (tensor): world matrix
+        scale_mat (tensor): scale matrix
+        invert (bool): whether to invert matrices (default: true)
+    '''
+    assert(pixels.shape[-1] == 2)
+    if scale_mat is None:
+        scale_mat = torch.eye(4).unsqueeze(0).repeat(
+            camera_mat.shape[0], 1, 1).to(camera_mat.device)
+
+    # Convert to pytorch
+    pixels, is_numpy = to_pytorch(pixels, True)
+    depth = to_pytorch(depth)
+    camera_mat = to_pytorch(camera_mat)
+    world_mat = to_pytorch(world_mat)
+    scale_mat = to_pytorch(scale_mat)
+
+    # Invert camera matrices
+    if invert:
+        camera_mat = torch.inverse(camera_mat)
+        world_mat = torch.inverse(world_mat)
+        scale_mat = torch.inverse(scale_mat)
+
+    # Transform pixels to homogen coordinates
+    pixels = pixels.permute(0, 2, 1)
+    pixels = torch.cat([pixels, torch.ones_like(pixels)], dim=1)
+
+    # Project pixels into camera space
+    if use_absolute_depth:
+        pixels[:, :2] = pixels[:, :2] * depth.permute(0, 2, 1).abs()
+        pixels[:, 2:3] = pixels[:, 2:3] * depth.permute(0, 2, 1)
+    else:
+        pixels[:, :3] = pixels[:, :3] * depth.permute(0, 2, 1)
+    
+    # Transform pixels to world space
+    p_world = scale_mat @ world_mat @ camera_mat @ pixels
+
+    # Transform p_world back to 3D coordinates
+    p_world = p_world[:, :3].permute(0, 2, 1)
+
+    if is_numpy:
+        p_world = p_world.numpy()
+    return p_world
+
+
+def transform_to_camera_space(p_world, world_mat, camera_mat=None, scale_mat=None):
+    ''' Transforms world points to camera space.
+        Args:
+        p_world (tensor): world points tensor of size B x N x 3
+        camera_mat (tensor): camera matrix
+        world_mat (tensor): world matrix
+        scale_mat (tensor): scale matrix
+    '''
+    batch_size, n_p, _ = p_world.shape
+    device = p_world.device
+
+    # Transform world points to homogen coordinates
+    p_world = torch.cat([p_world, torch.ones(
+        batch_size, n_p, 1).to(device)], dim=-1).permute(0, 2, 1)
+
+    # Apply matrices to transform p_world to camera space
+    if scale_mat is None:
+        if camera_mat is None:
+            p_cam = world_mat @ p_world
+        else:
+            p_cam = camera_mat @ world_mat @ p_world
+    else:
+        p_cam = camera_mat @ world_mat @ scale_mat @ p_world
+
+    # Transform points back to 3D coordinates
+    p_cam = p_cam[:, :3].permute(0, 2, 1)
+    return p_cam
+
+
+def origin_to_world(n_points, camera_mat, world_mat, scale_mat=None,
+                    invert=False):
+    ''' Transforms origin (camera location) to world coordinates.
+
+    Args:
+        n_points (int): how often the transformed origin is repeated in the
+            form (batch_size, n_points, 3)
+        camera_mat (tensor): camera matrix
+        world_mat (tensor): world matrix
+        scale_mat (tensor): scale matrix
+        invert (bool): whether to invert the matrices (default: true)
+    '''
+    batch_size = camera_mat.shape[0]
+    device = camera_mat.device
+    # Create origin in homogen coordinates
+    p = torch.zeros(batch_size, 4, n_points).to(device)
+    p[:, -1] = 1.
+
+    if scale_mat is None:
+        scale_mat = torch.eye(4).unsqueeze(
+            0).repeat(batch_size, 1, 1).to(device)
+
+    # Invert matrices
+    if invert:
+        camera_mat = torch.inverse(camera_mat)
+        world_mat = torch.inverse(world_mat)
+        scale_mat = torch.inverse(scale_mat)
+
+    # Apply transformation
+    p_world = scale_mat @ world_mat @ camera_mat @ p
+
+    # Transform points back to 3D coordinates
+    p_world = p_world[:, :3].permute(0, 2, 1)
+    return p_world
+
+
+def image_points_to_world(image_points, camera_mat, world_mat, scale_mat=None,
+                          invert=False, negative_depth=True):
+    ''' Transforms points on image plane to world coordinates.
+
+    In contrast to transform_to_world, no depth value is needed as points on
+    the image plane have a fixed depth of 1.
+
+    Args:
+        image_points (tensor): image points tensor of size B x N x 2
+        camera_mat (tensor): camera matrix
+        world_mat (tensor): world matrix
+        scale_mat (tensor): scale matrix
+        invert (bool): whether to invert matrices
+    '''
+    batch_size, n_pts, dim = image_points.shape
+    assert(dim == 2)
+    device = image_points.device
+    d_image = torch.ones(batch_size, n_pts, 1).to(device)
+    if negative_depth:
+        d_image *= -1.
+    return transform_to_world(image_points, d_image, camera_mat, world_mat,
+                              scale_mat, invert=invert)
+
+
+def image_points_to_camera(image_points, camera_mat, 
+                           invert=False, negative_depth=True, use_absolute_depth=True):
+    batch_size, n_pts, dim = image_points.shape
+    assert(dim == 2)
+    device = image_points.device
+    d_image = torch.ones(batch_size, n_pts, 1).to(device)
+    if negative_depth:
+        d_image *= -1.
+
+    # Convert to pytorch
+    pixels, is_numpy = to_pytorch(image_points, True)
+    depth = to_pytorch(d_image)
+    camera_mat = to_pytorch(camera_mat)
+
+    # Invert camera matrices
+    if invert:
+        camera_mat = torch.inverse(camera_mat)
+    
+    # Transform pixels to homogen coordinates
+    pixels = pixels.permute(0, 2, 1)
+    pixels = torch.cat([pixels, torch.ones_like(pixels)], dim=1)
+
+    # Project pixels into camera space
+    if use_absolute_depth:
+        pixels[:, :2] = pixels[:, :2] * depth.permute(0, 2, 1).abs()
+        pixels[:, 2:3] = pixels[:, 2:3] * depth.permute(0, 2, 1)
+    else:
+        pixels[:, :3] = pixels[:, :3] * depth.permute(0, 2, 1)
+
+    # Transform pixels to world space
+    p_camera = camera_mat @ pixels
+
+    # Transform p_world back to 3D coordinates
+    p_camera = p_camera[:, :3].permute(0, 2, 1)
+
+    if is_numpy:
+        p_camera = p_camera.numpy()
+    return p_camera
+
+
+def camera_points_to_image(camera_points, camera_mat, 
+                           invert=False, negative_depth=True, use_absolute_depth=True):
+    batch_size, n_pts, dim = camera_points.shape
+    assert(dim == 3)
+    device = camera_points.device
+
+    # Convert to pytorch
+    p_camera, is_numpy = to_pytorch(camera_points, True)
+    camera_mat = to_pytorch(camera_mat)
+
+    # Invert camera matrices
+    if invert:
+        camera_mat = torch.inverse(camera_mat)
+
+    # Transform world camera space to pixels
+    p_camera = p_camera.permute(0, 2, 1)  # B x 3 x N
+    pixels = camera_mat[:, :3, :3] @ p_camera
+
+    assert use_absolute_depth and negative_depth
+    pixels, p_depths = pixels[:, :2], pixels[:, 2:3]
+    p_depths = -p_depths  # negative depth
+    pixels = pixels / p_depths
+
+    pixels = pixels.permute(0, 2, 1)
+    if is_numpy:
+        pixels = pixels.numpy()
+    return pixels
+
+
+def angular_interpolation(res, camera_mat):
+    batch_size = camera_mat.shape[0]
+    device = camera_mat.device
+    input_rays  = image_points_to_camera(arange_pixels((res, res), batch_size, 
+        invert_y_axis=True).to(device), camera_mat)
+    output_rays = image_points_to_camera(arange_pixels((res * 2, res * 2), batch_size,
+        invert_y_axis=True).to(device), camera_mat)
+    input_rays  = input_rays / input_rays.norm(dim=-1, keepdim=True)
+    output_rays = output_rays / output_rays.norm(dim=-1, keepdim=True)
+
+    def dir2sph(v):
+        u = (v[..., :2] ** 2).sum(-1).sqrt()
+        theta = torch.atan2(u, v[..., 2]) / math.pi
+        phi = torch.atan2(v[..., 1], v[..., 0]) / math.pi
+        return torch.stack([theta, phi], 1)
+
+    input_rays  = dir2sph(input_rays).reshape(batch_size, 2, res, res)
+    output_rays = dir2sph(output_rays).reshape(batch_size, 2, res * 2, res * 2)
+    return input_rays
+
+
+def interpolate_sphere(z1, z2, t):
+    p = (z1 * z2).sum(dim=-1, keepdim=True)
+    p = p / z1.pow(2).sum(dim=-1, keepdim=True).sqrt()
+    p = p / z2.pow(2).sum(dim=-1, keepdim=True).sqrt()
+    omega = torch.acos(p)
+    s1 = torch.sin((1-t)*omega)/torch.sin(omega)
+    s2 = torch.sin(t*omega)/torch.sin(omega)
+    z = s1 * z1 + s2 * z2
+    return z
+
+
+def get_camera_rays(camera_matrices, pixels=None, res=None, margin=0):
+    device     = camera_matrices[0].device
+    batch_size = camera_matrices[0].shape[0]
+    if pixels is None:
+        assert res is not None
+        pixels = arange_pixels((res, res), batch_size, invert_y_axis=True, margin=margin).to(device)
+    n_points = pixels.size(1)
+    pixels_world = image_points_to_world(
+            pixels, camera_mat=camera_matrices[0],
+            world_mat=camera_matrices[1])
+    camera_world = origin_to_world(
+            n_points, camera_mat=camera_matrices[0],
+            world_mat=camera_matrices[1])
+    ray_vector = pixels_world - camera_world
+    ray_vector = ray_vector / ray_vector.norm(dim=-1, keepdim=True)
+    return ray_vector, camera_world, pixels_world
+
+
+def rotation_6d_to_matrix(d6: torch.Tensor) -> torch.Tensor:
+    """
+    Converts 6D rotation representation by Zhou et al. [1] to rotation matrix
+    using Gram--Schmidt orthogonalization per Section B of [1].
+    Args:
+        d6: 6D rotation representation, of size (*, 6)
+
+    Returns:
+        batch of rotation matrices of size (*, 3, 3)
+
+    [1] Zhou, Y., Barnes, C., Lu, J., Yang, J., & Li, H.
+    On the Continuity of Rotation Representations in Neural Networks.
+    IEEE Conference on Computer Vision and Pattern Recognition, 2019.
+    Retrieved from http://arxiv.org/abs/1812.07035
+    """
+
+    a1, a2 = d6[..., :3], d6[..., 3:]
+    b1 = F.normalize(a1, dim=-1)
+    b2 = a2 - (b1 * a2).sum(-1, keepdim=True) * b1
+    b2 = F.normalize(b2, dim=-1)
+    b3 = torch.cross(b1, b2, dim=-1)
+    return torch.stack((b1, b2, b3), dim=-2)
+
+
+def camera_9d_to_16d(d9):
+    d6, translation = d9[..., :6], d9[..., 6:]
+    rotation = rotation_6d_to_matrix(d6)
+    RT = torch.eye(4).to(device=d9.device, dtype=d9.dtype).reshape(
+        1, 4, 4).repeat(d6.size(0), 1, 1)
+    RT[:, :3, :3] = rotation
+    RT[:, :3, -1] = translation
+    return RT.reshape(-1, 16)
+
+def matrix_to_rotation_6d(matrix: torch.Tensor) -> torch.Tensor:
+    """
+    Converts rotation matrices to 6D rotation representation by Zhou et al. [1]
+    by dropping the last row. Note that 6D representation is not unique.
+    Args:
+        matrix: batch of rotation matrices of size (*, 3, 3)
+
+    Returns:
+        6D rotation representation, of size (*, 6)
+
+    [1] Zhou, Y., Barnes, C., Lu, J., Yang, J., & Li, H.
+    On the Continuity of Rotation Representations in Neural Networks.
+    IEEE Conference on Computer Vision and Pattern Recognition, 2019.
+    Retrieved from http://arxiv.org/abs/1812.07035
+    """
+    return matrix[..., :2, :].clone().reshape(*matrix.size()[:-2], 6)
+
+
+def depth2pts_outside(ray_o, ray_d, depth):
+    '''
+    ray_o, ray_d: [..., 3]
+    depth: [...]; inverse of distance to sphere origin
+    '''
+    # note: d1 becomes negative if this mid point is behind camera
+    d1 = -torch.sum(ray_d * ray_o, dim=-1) / torch.sum(ray_d * ray_d, dim=-1)
+    p_mid = ray_o + d1.unsqueeze(-1) * ray_d
+    p_mid_norm = torch.norm(p_mid, dim=-1)
+    ray_d_cos = 1. / torch.norm(ray_d, dim=-1)
+    d2 = torch.sqrt(1. - p_mid_norm * p_mid_norm) * ray_d_cos
+    p_sphere = ray_o + (d1 + d2).unsqueeze(-1) * ray_d
+
+    rot_axis = torch.cross(ray_o, p_sphere, dim=-1)
+    rot_axis = rot_axis / torch.norm(rot_axis, dim=-1, keepdim=True)
+    phi = torch.asin(p_mid_norm)
+    theta = torch.asin(p_mid_norm * depth)  # depth is inside [0, 1]
+    rot_angle = (phi - theta).unsqueeze(-1)     # [..., 1]
+
+    # now rotate p_sphere
+    # Rodrigues formula: https://en.wikipedia.org/wiki/Rodrigues%27_rotation_formula
+    p_sphere_new = p_sphere * torch.cos(rot_angle) + \
+                torch.cross(rot_axis, p_sphere, dim=-1) * torch.sin(rot_angle) + \
+                rot_axis * torch.sum(rot_axis*p_sphere, dim=-1, keepdim=True) * (1.-torch.cos(rot_angle))
+    p_sphere_new = p_sphere_new / torch.norm(p_sphere_new, dim=-1, keepdim=True)
+    pts = torch.cat((p_sphere_new, depth.unsqueeze(-1)), dim=-1)
+
+    # now calculate conventional depth
+    depth_real = 1. / (depth + TINY_NUMBER) * torch.cos(theta) * ray_d_cos + d1
+    return pts, depth_real
+
+
+def intersect_sphere(ray_o, ray_d, radius=1):
+    '''
+    ray_o, ray_d: [..., 3]
+    compute the depth of the intersection point between this ray and unit sphere
+    '''
+    # note: d1 becomes negative if this mid point is behind camera
+    d1 = -torch.sum(ray_d * ray_o, dim=-1) / torch.sum(ray_d * ray_d, dim=-1)
+    p = ray_o + d1.unsqueeze(-1) * ray_d
+    # consider the case where the ray does not intersect the sphere
+    ray_d_cos = 1. / torch.norm(ray_d, dim=-1)
+    d2 = radius ** 2 - torch.sum(p * p, dim=-1)
+    mask = (d2 > 0)
+    d2 = torch.sqrt(d2.clamp(min=1e-6)) * ray_d_cos
+    d1, d2 = d1.unsqueeze(-1), d2.unsqueeze(-1)
+    depth_range = [d1 - d2, d1 + d2]
+    return depth_range, mask
+
+
+def normalize(x, axis=-1, order=2):
+    if isinstance(x, torch.Tensor):
+        l2 = x.norm(p=order, dim=axis, keepdim=True)
+        return x / (l2 + 1e-8), l2
+
+    else:
+        l2 = np.linalg.norm(x, order, axis)
+        l2 = np.expand_dims(l2, axis)
+        l2[l2==0] = 1
+        return x / l2, l2
+
+
+def sample_pdf(bins, weights, N_importance, det=False, eps=1e-5):
+    """
+    Sample @N_importance samples from @bins with distribution defined by @weights.
+    Inputs:
+        bins: (N_rays, N_samples_+1) where N_samples_ is "the number of coarse samples per ray - 2"
+        weights: (N_rays, N_samples_)
+        N_importance: the number of samples to draw from the distribution
+        det: deterministic or not
+        eps: a small number to prevent division by zero
+    Outputs:
+        samples: the sampled samples
+    Source: https://github.com/kwea123/nerf_pl/blob/master/models/rendering.py
+    """
+    N_rays, N_samples_ = weights.shape
+    weights = weights + eps # prevent division by zero (don't do inplace op!)
+    pdf = weights / torch.sum(weights, -1, keepdim=True) # (N_rays, N_samples_)
+    cdf = torch.cumsum(pdf, -1) # (N_rays, N_samples), cumulative distribution function
+    cdf = torch.cat([torch.zeros_like(cdf[: ,:1]), cdf], -1)  # (N_rays, N_samples_+1)
+                                                               # padded to 0~1 inclusive
+
+    if det:
+        u = torch.linspace(0, 1, N_importance, device=bins.device)
+        u = u.expand(N_rays, N_importance)
+    else:
+        u = torch.rand(N_rays, N_importance, device=bins.device)
+    u = u.contiguous()
+
+    inds = torch.searchsorted(cdf, u)
+    below = torch.clamp_min(inds-1, 0)
+    above = torch.clamp_max(inds, N_samples_)
+
+    inds_sampled = torch.stack([below, above], -1).view(N_rays, 2*N_importance)
+    cdf_g = torch.gather(cdf, 1, inds_sampled)
+    cdf_g = cdf_g.view(N_rays, N_importance, 2)
+    bins_g = torch.gather(bins, 1, inds_sampled).view(N_rays, N_importance, 2)
+
+    denom = cdf_g[...,1]-cdf_g[...,0]
+    denom[denom<eps] = 1 # denom equals 0 means a bin has weight 0, in which case it will not be sampled
+                         # anyway, therefore any value for it is fine (set to 1 here)
+
+    samples = bins_g[...,0] + (u-cdf_g[...,0])/denom * (bins_g[...,1]-bins_g[...,0])
+    return samples
+
+
+def normalization_inverse_sqrt_dist_centered(x_in_world, view_cell_center, max_depth):
+    localized = x_in_world - view_cell_center
+    local = torch.sqrt(torch.linalg.norm(localized, dim=-1))
+    res = localized / (math.sqrt(max_depth) * local[..., None])
+    return res
+
+
+######################################################################################

dnnlib/filters.py

@@ -0,0 +1,81 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+
+import math
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+def kaiser_attenuation(n_taps, f_h, sr):
+    df = (2 * f_h) / (sr / 2)
+    return 2.285 * (n_taps - 1) * math.pi * df + 7.95
+
+
+def kaiser_beta(n_taps, f_h, sr):
+    atten = kaiser_attenuation(n_taps, f_h, sr)
+
+    if atten > 50:
+        return 0.1102 * (atten - 8.7)
+
+    elif 50 >= atten >= 21:
+        return 0.5842 * (atten - 21) ** 0.4 + 0.07886 * (atten - 21)
+
+    else:
+        return 0.0
+
+def sinc(x, eps=1e-10):
+    y = torch.sin(math.pi * x) / (math.pi * x + eps)
+    y = y.masked_fill(x.eq(0), 1.0)
+    return y
+
+
+def kaiser_window(n_taps, f_h, sr):
+    beta = kaiser_beta(n_taps, f_h, sr)
+    ind = torch.arange(n_taps) - (n_taps - 1) / 2
+    return torch.i0(beta * torch.sqrt(1 - ((2 * ind) / (n_taps - 1)) ** 2)) / torch.i0(
+        torch.tensor(beta)
+    )
+
+
+def lowpass_filter(n_taps, cutoff, band_half, sr):
+    window = kaiser_window(n_taps, band_half, sr)
+    ind = torch.arange(n_taps) - (n_taps - 1) / 2
+    lowpass = 2 * cutoff / sr * sinc(2 * cutoff / sr * ind) * window
+    return lowpass
+
+
+def filter_parameters(
+    n_layer,
+    n_critical,
+    sr_max,
+    cutoff_0,
+    cutoff_n,
+    stopband_0,
+    stopband_n
+):
+    cutoffs = []
+    stopbands = []
+    srs = []
+    band_halfs = []
+
+    for i in range(n_layer):
+        f_c = cutoff_0 * (cutoff_n / cutoff_0) ** min(i / (n_layer - n_critical), 1)
+        f_t = stopband_0 * (stopband_n / stopband_0) ** min(
+            i / (n_layer - n_critical), 1
+        )
+        s_i = 2 ** math.ceil(math.log(min(2 * f_t, sr_max), 2))
+        f_h = max(f_t, s_i / 2) - f_c
+
+        cutoffs.append(f_c)
+        stopbands.append(f_t)
+        srs.append(s_i)
+        band_halfs.append(f_h)
+
+    return {
+        "cutoffs": cutoffs,
+        "stopbands": stopbands,
+        "srs": srs,
+        "band_halfs": band_halfs,
+    }

dnnlib/geometry.py

@@ -0,0 +1,406 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+
+import torch
+import torch.nn.functional as F
+import math
+import random
+import numpy as np
+
+
+def positional_encoding(p, size, pe='normal', use_pos=False):
+    if pe == 'gauss':
+        p_transformed = np.pi * p @ size
+        p_transformed = torch.cat(
+            [torch.sin(p_transformed), torch.cos(p_transformed)], dim=-1)
+    else:
+        p_transformed = torch.cat([torch.cat(
+            [torch.sin((2 ** i) * np.pi * p),
+            torch.cos((2 ** i) * np.pi * p)],
+            dim=-1) for i in range(size)], dim=-1)
+    if use_pos:
+        p_transformed = torch.cat([p_transformed, p], -1)
+    return p_transformed
+
+
+def upsample(img_nerf, size, filter=None):
+    up = size // img_nerf.size(-1)
+    if up <= 1:
+        return img_nerf
+
+    if filter is not None:
+        from torch_utils.ops import upfirdn2d
+        for _ in range(int(math.log2(up))):
+            img_nerf = upfirdn2d.downsample2d(img_nerf, filter, up=2)
+    else:
+        img_nerf = F.interpolate(img_nerf, (size, size), mode='bilinear', align_corners=False)
+    return img_nerf
+
+
+def downsample(img0, size, filter=None):
+    down = img0.size(-1) // size
+    if down <= 1:
+        return img0
+
+    if filter is not None:    
+        from torch_utils.ops import upfirdn2d
+        for _ in range(int(math.log2(down))):
+            img0 = upfirdn2d.downsample2d(img0, filter, down=2)
+    else:
+        img0 = F.interpolate(img0, (size, size), mode='bilinear', align_corners=False)
+    return img0
+        
+
+def normalize_vecs(vectors: torch.Tensor) -> torch.Tensor:
+    """
+    Normalize vector lengths.
+    """
+    return vectors / (torch.norm(vectors, dim=-1, keepdim=True))
+
+
+def repeat_vecs(vecs, n, dim=0):
+    return torch.stack(n*[vecs], dim=dim)
+
+
+def get_grids(H, W, device, align=True):
+    ch = 1 if align else 1 - (1 / H)
+    cw = 1 if align else 1 - (1 / W)
+    x, y = torch.meshgrid(torch.linspace(-cw, cw, W, device=device),
+                          torch.linspace(ch, -ch, H, device=device))
+    return torch.stack([x, y], -1)
+
+
+def local_ensemble(pi, po, resolution):
+    ii = range(resolution)
+    ia = torch.tensor([max((i - 1)//2, 0) for i in ii]).long()
+    ib = torch.tensor([min((i + 1)//2, resolution//2-1) for i in ii]).long()
+    
+    ul = torch.meshgrid(ia, ia)
+    ur = torch.meshgrid(ia, ib)
+    ll = torch.meshgrid(ib, ia)
+    lr = torch.meshgrid(ib, ib)
+    
+    d_ul, p_ul = po - pi[ul], torch.stack(ul, -1)
+    d_ur, p_ur = po - pi[ur], torch.stack(ur, -1)
+    d_ll, p_ll = po - pi[ll], torch.stack(ll, -1)
+    d_lr, p_lr = po - pi[lr], torch.stack(lr, -1)
+    
+    c_ul = d_ul.prod(dim=-1).abs()
+    c_ur = d_ur.prod(dim=-1).abs()
+    c_ll = d_ll.prod(dim=-1).abs()
+    c_lr = d_lr.prod(dim=-1).abs()
+
+    D = torch.stack([d_ul, d_ur, d_ll, d_lr], 0)
+    P = torch.stack([p_ul, p_ur, p_ll, p_lr], 0)
+    C = torch.stack([c_ul, c_ur, c_ll, c_lr], 0)
+    C = C / C.sum(dim=0, keepdim=True)
+    return D, P, C
+
+
+def get_initial_rays_trig(num_steps, fov, resolution, ray_start, ray_end, device='cpu'):
+    """Returns sample points, z_vals, ray directions in camera space."""
+
+    W, H = resolution
+    # Create full screen NDC (-1 to +1) coords [x, y, 0, 1].
+    # Y is flipped to follow image memory layouts.
+    x, y = torch.meshgrid(torch.linspace(-1, 1, W, device=device),
+                          torch.linspace(1, -1, H, device=device))
+    x = x.T.flatten()
+    y = y.T.flatten()
+    z = -torch.ones_like(x, device=device) / math.tan((2 * math.pi * fov / 360)/2)
+
+    rays_d_cam = normalize_vecs(torch.stack([x, y, z], -1))
+
+    z_vals = torch.linspace(ray_start, ray_end, num_steps, device=device).reshape(1, num_steps, 1).repeat(W*H, 1, 1)
+    points = rays_d_cam.unsqueeze(1).repeat(1, num_steps, 1) * z_vals
+    return points, z_vals, rays_d_cam
+
+
+def sample_camera_positions(
+    device, n=1, r=1, horizontal_stddev=1, vertical_stddev=1, 
+    horizontal_mean=math.pi*0.5, vertical_mean=math.pi*0.5, mode='normal'):
+    """
+    Samples n random locations along a sphere of radius r. 
+    Uses a gaussian distribution for pitch and yaw
+    """
+    if mode == 'uniform':
+        theta = (torch.rand((n, 1),device=device) - 0.5) * 2 * horizontal_stddev + horizontal_mean
+        phi = (torch.rand((n, 1),device=device) - 0.5) * 2 * vertical_stddev + vertical_mean
+
+    elif mode == 'normal' or mode == 'gaussian':
+        theta = torch.randn((n, 1), device=device) * horizontal_stddev + horizontal_mean
+        phi = torch.randn((n, 1), device=device) * vertical_stddev + vertical_mean
+
+    elif mode == 'hybrid':
+        if random.random() < 0.5:
+            theta = (torch.rand((n, 1),device=device) - 0.5) * 2 * horizontal_stddev * 2 + horizontal_mean
+            phi = (torch.rand((n, 1),device=device) - 0.5) * 2 * vertical_stddev * 2 + vertical_mean
+        else:
+            theta = torch.randn((n, 1), device=device) * horizontal_stddev + horizontal_mean
+            phi = torch.randn((n, 1), device=device) * vertical_stddev + vertical_mean
+    else:
+        phi = torch.ones((n, 1), device=device, dtype=torch.float) * vertical_mean
+        theta = torch.ones((n, 1), device=device, dtype=torch.float) * horizontal_mean
+
+    phi = torch.clamp(phi, 1e-5, math.pi - 1e-5)
+
+    output_points = torch.zeros((n, 3), device=device)# torch.cuda.FloatTensor(n, 3).fill_(0)#torch.zeros((n, 3))
+
+    output_points[:, 0:1] = r*torch.sin(phi) * torch.cos(theta)
+    output_points[:, 2:3] = r*torch.sin(phi) * torch.sin(theta)
+    output_points[:, 1:2] = r*torch.cos(phi)
+
+    return output_points, phi, theta
+
+
+def perturb_points(points, z_vals, ray_directions, device):
+    distance_between_points = z_vals[:,:,1:2,:] - z_vals[:,:,0:1,:]
+    offset = (torch.rand(z_vals.shape, device=device)-0.5) * distance_between_points
+    z_vals = z_vals + offset
+    points = points + offset * ray_directions.unsqueeze(2)
+    return points, z_vals
+
+
+def create_cam2world_matrix(forward_vector, origin, device=None):
+    """Takes in the direction the camera is pointing and the camera origin and returns a world2cam matrix."""
+
+    forward_vector = normalize_vecs(forward_vector)
+    up_vector = torch.tensor([0, 1, 0], dtype=torch.float, device=device).expand_as(forward_vector)
+    left_vector = normalize_vecs(torch.cross(up_vector, forward_vector, dim=-1))
+    up_vector = normalize_vecs(torch.cross(forward_vector, left_vector, dim=-1))
+
+    rotation_matrix = torch.eye(4, device=device).unsqueeze(0).repeat(forward_vector.shape[0], 1, 1)
+    rotation_matrix[:, :3, :3] = torch.stack((-left_vector, up_vector, -forward_vector), axis=-1)
+
+    translation_matrix = torch.eye(4, device=device).unsqueeze(0).repeat(forward_vector.shape[0], 1, 1)
+    translation_matrix[:, :3, 3] = origin
+
+    cam2world = translation_matrix @ rotation_matrix
+
+    return cam2world
+
+
+def transform_sampled_points(
+    points, z_vals, ray_directions, device, 
+    h_stddev=1, v_stddev=1, h_mean=math.pi * 0.5, 
+    v_mean=math.pi * 0.5, mode='normal'):
+    """
+    points: batch_size x total_pixels x num_steps x 3
+    z_vals: batch_size x total_pixels x num_steps
+    """
+    n, num_rays, num_steps, channels = points.shape
+    points, z_vals = perturb_points(points, z_vals, ray_directions, device)
+    camera_origin, pitch, yaw = sample_camera_positions(
+        n=points.shape[0], r=1, 
+        horizontal_stddev=h_stddev, vertical_stddev=v_stddev, 
+        horizontal_mean=h_mean, vertical_mean=v_mean, 
+        device=device, mode=mode)
+    forward_vector = normalize_vecs(-camera_origin)
+    cam2world_matrix = create_cam2world_matrix(forward_vector, camera_origin, device=device)
+
+    points_homogeneous = torch.ones((points.shape[0], points.shape[1], points.shape[2], points.shape[3] + 1), device=device)
+    points_homogeneous[:, :, :, :3] = points
+
+    # should be n x 4 x 4 , n x r^2 x num_steps x 4
+    transformed_points = torch.bmm(cam2world_matrix, points_homogeneous.reshape(n, -1, 4).permute(0,2,1)).permute(0, 2, 1).reshape(n, num_rays, num_steps, 4)
+    transformed_ray_directions = torch.bmm(cam2world_matrix[..., :3, :3], ray_directions.reshape(n, -1, 3).permute(0,2,1)).permute(0, 2, 1).reshape(n, num_rays, 3)
+
+    homogeneous_origins = torch.zeros((n, 4, num_rays), device=device)
+    homogeneous_origins[:, 3, :] = 1
+
+    transformed_ray_origins = torch.bmm(cam2world_matrix, homogeneous_origins).permute(0, 2, 1).reshape(n, num_rays, 4)[..., :3]
+    return transformed_points[..., :3], z_vals, transformed_ray_directions, transformed_ray_origins, pitch, yaw
+
+
+def integration(
+    rgb_sigma, z_vals, device, noise_std=0.5,
+    last_back=False, white_back=False, clamp_mode=None, fill_mode=None):
+
+    rgbs = rgb_sigma[..., :3]
+    sigmas = rgb_sigma[..., 3:]
+
+    deltas = z_vals[..., 1:, :] - z_vals[..., :-1, :]
+    delta_inf = 1e10 * torch.ones_like(deltas[..., :1, :])
+    deltas = torch.cat([deltas, delta_inf], -2)
+
+    if noise_std > 0:
+        noise = torch.randn(sigmas.shape, device=device) * noise_std
+    else:
+        noise = 0
+
+    if clamp_mode == 'softplus':
+        alphas = 1 - torch.exp(-deltas * (F.softplus(sigmas + noise)))
+    elif clamp_mode == 'relu':
+        alphas = 1 - torch.exp(-deltas * (F.relu(sigmas + noise)))
+    else:
+        raise "Need to choose clamp mode"
+
+    alphas_shifted = torch.cat([torch.ones_like(alphas[..., :1, :]), 1-alphas + 1e-10], -2)
+    weights = alphas * torch.cumprod(alphas_shifted, -2)[..., :-1, :]
+    weights_sum = weights.sum(-2)
+
+    if last_back:
+        weights[..., -1, :] += (1 - weights_sum)
+
+    rgb_final = torch.sum(weights * rgbs, -2)
+    depth_final = torch.sum(weights * z_vals, -2)
+
+    if white_back:
+        rgb_final = rgb_final + 1-weights_sum
+
+    if fill_mode == 'debug':
+        rgb_final[weights_sum.squeeze(-1) < 0.9] = torch.tensor([1., 0, 0], device=rgb_final.device)
+    elif fill_mode == 'weight':
+        rgb_final = weights_sum.expand_as(rgb_final)
+
+    return rgb_final, depth_final, weights
+
+
+def get_sigma_field_np(nerf, styles, resolution=512, block_resolution=64):
+    # return numpy array of forwarded sigma value
+    bound = (nerf.depth_range[1] - nerf.depth_range[0]) * 0.5
+    X = torch.linspace(-bound, bound, resolution).split(block_resolution)
+
+    sigma_np = np.zeros([resolution, resolution, resolution], dtype=np.float32)
+
+    for xi, xs in enumerate(X):
+        for yi, ys in enumerate(X):
+            for zi, zs in enumerate(X):
+                xx, yy, zz = torch.meshgrid(xs, ys, zs)
+                pts = torch.stack([xx, yy, zz], dim=-1).unsqueeze(0).to(styles.device)  # B, H, H, H, C
+                block_shape = [1, len(xs), len(ys), len(zs)]
+                feat_out, sigma_out = nerf.fg_nerf.forward_style2(pts, None, block_shape, ws=styles)
+                sigma_np[xi * block_resolution: xi * block_resolution + len(xs), \
+                yi * block_resolution: yi * block_resolution + len(ys), \
+                zi * block_resolution: zi * block_resolution + len(zs)] = sigma_out.reshape(block_shape[1:]).detach().cpu().numpy()
+
+    return sigma_np, bound
+
+
+def extract_geometry(nerf, styles, resolution, threshold):
+    import mcubes
+
+    print('threshold: {}'.format(threshold))
+    u, bound = get_sigma_field_np(nerf, styles, resolution)
+    vertices, triangles = mcubes.marching_cubes(u, threshold)
+    b_min_np = np.array([-bound, -bound, -bound])
+    b_max_np = np.array([ bound,  bound,  bound])
+
+    vertices = vertices / (resolution - 1.0) * (b_max_np - b_min_np)[None, :] + b_min_np[None, :]
+    return vertices.astype('float32'), triangles
+
+
+def render_mesh(meshes, camera_matrices, render_noise=True):
+    from pytorch3d.renderer import (
+        FoVPerspectiveCameras, look_at_view_transform,
+        RasterizationSettings, BlendParams,
+        MeshRenderer, MeshRasterizer, HardPhongShader, TexturesVertex
+    )
+    from pytorch3d.ops import interpolate_face_attributes
+    from pytorch3d.structures.meshes import Meshes
+
+    intrinsics, poses, _, _ = camera_matrices
+    device = poses.device
+    c2w = torch.matmul(poses, torch.diag(torch.tensor([-1.0, 1.0, -1.0, 1.0], device=device))[None, :, :])  # Different camera model...
+    w2c = torch.inverse(c2w)
+    R = c2w[:, :3, :3]
+    T = w2c[:, :3, 3]   # So weird..... Why one is c2w and another is w2c?
+    focal = intrinsics[0, 0, 0]
+    fov = torch.arctan(focal) * 2.0 / np.pi * 180
+    
+
+    colors = []
+    offset = 1
+    for res, (mesh, face_vert_noise) in meshes.items():
+        raster_settings = RasterizationSettings(
+                image_size=res,
+                blur_radius=0.0,
+                faces_per_pixel=1,
+        )
+        mesh = Meshes(
+            verts=[torch.from_numpy(mesh.vertices).float().to(device)],
+            faces=[torch.from_numpy(mesh.faces).long().to(device)])
+
+        _colors = []
+        for i in range(len(poses)):
+            cameras = FoVPerspectiveCameras(device=device, R=R[i: i+1], T=T[i: i+1], fov=fov)
+            rasterizer = MeshRasterizer(cameras=cameras, raster_settings=raster_settings)
+            pix_to_face, zbuf, bary_coord, dists = rasterizer(mesh)
+            color = interpolate_face_attributes(pix_to_face, bary_coord, face_vert_noise).squeeze()
+
+            # hack
+            color[offset:, offset:] = color[:-offset, :-offset]
+            _colors += [color]
+        color = torch.stack(_colors, 0).permute(0,3,1,2)
+        colors += [color]
+        offset *= 2
+    return colors
+
+
+def rotate_vects(v, theta):
+    theta = theta / math.pi * 2
+    theta = theta + (theta < 0).type_as(theta) * 4
+    v  = v.reshape(v.size(0), v.size(1) // 4, 4, v.size(2), v.size(3))
+    vs = []
+    order  = [0,2,3,1]  # Not working
+    iorder = [0,3,1,2]  # Not working
+    for b in range(len(v)):
+        if (theta[b] - 0) < 1e-6:
+            u, l = 0, 0
+        elif (theta[b] - 1) < 1e-6:
+            u, l = 0, 1
+        elif (theta[b] - 2) < 1e-6:
+            u, l = 0, 2
+        elif (theta[b] - 3) < 1e-6:
+            u, l = 0, 3
+        else:
+            u, l = math.modf(theta[b])
+        l, r = int(l), int(l + 1) % 4
+        vv = v[b, :, order]  # 0 -> 1 -> 3 -> 2
+        vl   = torch.cat([vv[:, l:], vv[:, :l]], 1)
+        if u > 0:
+            vr = torch.cat([vv[:, r:], vv[:, :r]], 1)
+            vv = vl * (1-u) + vr * u
+        else:
+            vv = vl
+        vs.append(vv[:, iorder])
+    v = torch.stack(vs, 0)
+    v = v.reshape(v.size(0), -1, v.size(-2), v.size(-1))
+    return v
+
+
+def generate_option_outputs(render_option):
+    # output debugging outputs (not used in normal rendering process)
+    if ('depth' in render_option.split(',')):    
+        img = camera_world[:, :1] + fg_depth_map * ray_vector
+        img = reformat(img, tgt_res)
+
+        if 'gradient' in render_option.split(','):
+            points = (camera_world[:,:,None]+di[:,:,:,None]*ray_vector[:,:,None]).reshape(
+                batch_size, tgt_res, tgt_res, di.size(-1), 3)
+            with torch.enable_grad():
+                gradients = self.fg_nerf.forward_style2(
+                    points, None, [batch_size, tgt_res, di.size(-1), tgt_res], get_normal=True,
+                    ws=styles, z_shape=z_shape_obj, z_app=z_app_obj).reshape(
+                        batch_size, di.size(-1), 3, tgt_res * tgt_res).permute(0,3,1,2)
+                avg_grads = (gradients * fg_weights.unsqueeze(-1)).sum(-2)
+            normal = reformat(normalize(avg_grads, axis=2)[0], tgt_res)
+            img = normal
+
+        if 'value' in render_option.split(','):
+            fg_feat = fg_feat[:,:,3:].norm(dim=-1,keepdim=True)
+            img = reformat(fg_feat.repeat(1,1,3), tgt_res) / fg_feat.max() * 2 - 1
+            
+        if 'opacity' in render_option.split(','):
+            opacity = bg_lambda.unsqueeze(-1).repeat(1,1,3) * 2 - 1
+            img = reformat(opacity, tgt_res)
+
+        if 'normal' in render_option.split(','):
+            shift_l, shift_r = img[:,:,2:,:], img[:,:,:-2,:]
+            shift_u, shift_d = img[:,:,:,2:], img[:,:,:,:-2]
+            diff_hor = normalize(shift_r - shift_l, axis=1)[0][:, :, :, 1:-1]
+            diff_ver = normalize(shift_u - shift_d, axis=1)[0][:, :, 1:-1, :]
+            normal = torch.cross(diff_hor, diff_ver, dim=1)
+            img = normalize(normal, axis=1)[0]
+        
+        return {'full_out': (None, img), 'reg_loss': {}}

dnnlib/util.py

@@ -0,0 +1,531 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+"""Miscellaneous utility classes and functions."""
+
+import ctypes
+import fnmatch
+import importlib
+import inspect
+import numpy as np
+import os
+import shutil
+import sys
+import types
+import io
+import pickle
+import re
+import requests
+import html
+import hashlib
+import glob
+import tempfile
+import urllib
+import urllib.request
+import uuid
+import torch
+
+from distutils.util import strtobool
+from typing import Any, List, Tuple, Union
+
+
+# Util classes
+# ------------------------------------------------------------------------------------------
+
+
+class EasyDict(dict):
+    """Convenience class that behaves like a dict but allows access with the attribute syntax."""
+
+    def __getattr__(self, name: str) -> Any:
+        try:
+            return self[name]
+        except KeyError:
+            raise AttributeError(name)
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        self[name] = value
+
+    def __delattr__(self, name: str) -> None:
+        del self[name]
+
+
+class Logger(object):
+    """Redirect stderr to stdout, optionally print stdout to a file, and optionally force flushing on both stdout and the file."""
+
+    def __init__(self, file_name: str = None, file_mode: str = "w", should_flush: bool = True):
+        self.file = None
+
+        if file_name is not None:
+            self.file = open(file_name, file_mode)
+
+        self.should_flush = should_flush
+        self.stdout = sys.stdout
+        self.stderr = sys.stderr
+
+        sys.stdout = self
+        sys.stderr = self
+
+    def __enter__(self) -> "Logger":
+        return self
+
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
+        self.close()
+
+    def write(self, text: Union[str, bytes]) -> None:
+        """Write text to stdout (and a file) and optionally flush."""
+        if isinstance(text, bytes):
+            text = text.decode()
+        if len(text) == 0: # workaround for a bug in VSCode debugger: sys.stdout.write(''); sys.stdout.flush() => crash
+            return
+
+        if self.file is not None:
+            self.file.write(text)
+
+        self.stdout.write(text)
+
+        if self.should_flush:
+            self.flush()
+
+    def flush(self) -> None:
+        """Flush written text to both stdout and a file, if open."""
+        if self.file is not None:
+            self.file.flush()
+
+        self.stdout.flush()
+
+    def close(self) -> None:
+        """Flush, close possible files, and remove stdout/stderr mirroring."""
+        self.flush()
+
+        # if using multiple loggers, prevent closing in wrong order
+        if sys.stdout is self:
+            sys.stdout = self.stdout
+        if sys.stderr is self:
+            sys.stderr = self.stderr
+
+        if self.file is not None:
+            self.file.close()
+            self.file = None
+
+
+# Cache directories
+# ------------------------------------------------------------------------------------------
+
+_dnnlib_cache_dir = None
+
+def set_cache_dir(path: str) -> None:
+    global _dnnlib_cache_dir
+    _dnnlib_cache_dir = path
+
+def make_cache_dir_path(*paths: str) -> str:
+    if _dnnlib_cache_dir is not None:
+        return os.path.join(_dnnlib_cache_dir, *paths)
+    if 'DNNLIB_CACHE_DIR' in os.environ:
+        return os.path.join(os.environ['DNNLIB_CACHE_DIR'], *paths)
+    if 'HOME' in os.environ:
+        return os.path.join(os.environ['HOME'], '.cache', 'dnnlib', *paths)
+    if 'USERPROFILE' in os.environ:
+        return os.path.join(os.environ['USERPROFILE'], '.cache', 'dnnlib', *paths)
+    return os.path.join(tempfile.gettempdir(), '.cache', 'dnnlib', *paths)
+
+# Small util functions
+# ------------------------------------------------------------------------------------------
+
+
+def format_time(seconds: Union[int, float]) -> str:
+    """Convert the seconds to human readable string with days, hours, minutes and seconds."""
+    s = int(np.rint(seconds))
+
+    if s < 60:
+        return "{0}s".format(s)
+    elif s < 60 * 60:
+        return "{0}m {1:02}s".format(s // 60, s % 60)
+    elif s < 24 * 60 * 60:
+        return "{0}h {1:02}m {2:02}s".format(s // (60 * 60), (s // 60) % 60, s % 60)
+    else:
+        return "{0}d {1:02}h {2:02}m".format(s // (24 * 60 * 60), (s // (60 * 60)) % 24, (s // 60) % 60)
+
+
+def ask_yes_no(question: str) -> bool:
+    """Ask the user the question until the user inputs a valid answer."""
+    while True:
+        try:
+            print("{0} [y/n]".format(question))
+            return strtobool(input().lower())
+        except ValueError:
+            pass
+
+
+def tuple_product(t: Tuple) -> Any:
+    """Calculate the product of the tuple elements."""
+    result = 1
+
+    for v in t:
+        result *= v
+
+    return result
+
+
+_str_to_ctype = {
+    "uint8": ctypes.c_ubyte,
+    "uint16": ctypes.c_uint16,
+    "uint32": ctypes.c_uint32,
+    "uint64": ctypes.c_uint64,
+    "int8": ctypes.c_byte,
+    "int16": ctypes.c_int16,
+    "int32": ctypes.c_int32,
+    "int64": ctypes.c_int64,
+    "float32": ctypes.c_float,
+    "float64": ctypes.c_double
+}
+
+
+def get_dtype_and_ctype(type_obj: Any) -> Tuple[np.dtype, Any]:
+    """Given a type name string (or an object having a __name__ attribute), return matching Numpy and ctypes types that have the same size in bytes."""
+    type_str = None
+
+    if isinstance(type_obj, str):
+        type_str = type_obj
+    elif hasattr(type_obj, "__name__"):
+        type_str = type_obj.__name__
+    elif hasattr(type_obj, "name"):
+        type_str = type_obj.name
+    else:
+        raise RuntimeError("Cannot infer type name from input")
+
+    assert type_str in _str_to_ctype.keys()
+
+    my_dtype = np.dtype(type_str)
+    my_ctype = _str_to_ctype[type_str]
+
+    assert my_dtype.itemsize == ctypes.sizeof(my_ctype)
+
+    return my_dtype, my_ctype
+
+
+def is_pickleable(obj: Any) -> bool:
+    try:
+        with io.BytesIO() as stream:
+            pickle.dump(obj, stream)
+        return True
+    except:
+        return False
+
+
+# Functionality to import modules/objects by name, and call functions by name
+# ------------------------------------------------------------------------------------------
+
+def get_module_from_obj_name(obj_name: str) -> Tuple[types.ModuleType, str]:
+    """Searches for the underlying module behind the name to some python object.
+    Returns the module and the object name (original name with module part removed)."""
+
+    # allow convenience shorthands, substitute them by full names
+    obj_name = re.sub("^np.", "numpy.", obj_name)
+    obj_name = re.sub("^tf.", "tensorflow.", obj_name)
+
+    # list alternatives for (module_name, local_obj_name)
+    parts = obj_name.split(".")
+    name_pairs = [(".".join(parts[:i]), ".".join(parts[i:])) for i in range(len(parts), 0, -1)]
+
+    # try each alternative in turn
+    for module_name, local_obj_name in name_pairs:
+        try:
+            module = importlib.import_module(module_name) # may raise ImportError
+            get_obj_from_module(module, local_obj_name) # may raise AttributeError
+            return module, local_obj_name
+        except:
+            pass
+
+    # maybe some of the modules themselves contain errors?
+    for module_name, _local_obj_name in name_pairs:
+        try:
+            importlib.import_module(module_name) # may raise ImportError
+        except ImportError:
+            if not str(sys.exc_info()[1]).startswith("No module named '" + module_name + "'"):
+                raise
+
+    # maybe the requested attribute is missing?
+    for module_name, local_obj_name in name_pairs:
+        try:
+            module = importlib.import_module(module_name) # may raise ImportError
+            get_obj_from_module(module, local_obj_name) # may raise AttributeError
+        except ImportError:
+            pass
+
+    # we are out of luck, but we have no idea why
+    raise ImportError(obj_name)
+
+
+def get_obj_from_module(module: types.ModuleType, obj_name: str) -> Any:
+    """Traverses the object name and returns the last (rightmost) python object."""
+    if obj_name == '':
+        return module
+    obj = module
+    for part in obj_name.split("."):
+        obj = getattr(obj, part)
+    return obj
+
+
+def get_obj_by_name(name: str) -> Any:
+    """Finds the python object with the given name."""
+    module, obj_name = get_module_from_obj_name(name)
+    return get_obj_from_module(module, obj_name)
+
+
+def call_func_by_name(*args, func_name: str = None, **kwargs) -> Any:
+    """Finds the python object with the given name and calls it as a function."""
+    assert func_name is not None
+    func_obj = get_obj_by_name(func_name)
+    assert callable(func_obj)
+    return func_obj(*args, **kwargs)
+
+
+def construct_class_by_name(*args, class_name: str = None, **kwargs) -> Any:
+    """Finds the python class with the given name and constructs it with the given arguments."""
+    return call_func_by_name(*args, func_name=class_name, **kwargs)
+
+
+def get_module_dir_by_obj_name(obj_name: str) -> str:
+    """Get the directory path of the module containing the given object name."""
+    module, _ = get_module_from_obj_name(obj_name)
+    return os.path.dirname(inspect.getfile(module))
+
+
+def is_top_level_function(obj: Any) -> bool:
+    """Determine whether the given object is a top-level function, i.e., defined at module scope using 'def'."""
+    return callable(obj) and obj.__name__ in sys.modules[obj.__module__].__dict__
+
+
+def get_top_level_function_name(obj: Any) -> str:
+    """Return the fully-qualified name of a top-level function."""
+    assert is_top_level_function(obj)
+    module = obj.__module__
+    if module == '__main__':
+        module = os.path.splitext(os.path.basename(sys.modules[module].__file__))[0]
+    return module + "." + obj.__name__
+
+
+# File system helpers
+# ------------------------------------------------------------------------------------------
+
+def list_dir_recursively_with_ignore(dir_path: str, ignores: List[str] = None, add_base_to_relative: bool = False) -> List[Tuple[str, str]]:
+    """List all files recursively in a given directory while ignoring given file and directory names.
+    Returns list of tuples containing both absolute and relative paths."""
+    assert os.path.isdir(dir_path)
+    base_name = os.path.basename(os.path.normpath(dir_path))
+
+    if ignores is None:
+        ignores = []
+
+    result = []
+
+    for root, dirs, files in os.walk(dir_path, topdown=True):
+        for ignore_ in ignores:
+            dirs_to_remove = [d for d in dirs if fnmatch.fnmatch(d, ignore_)]
+
+            # dirs need to be edited in-place
+            for d in dirs_to_remove:
+                dirs.remove(d)
+
+            files = [f for f in files if not fnmatch.fnmatch(f, ignore_)]
+
+        absolute_paths = [os.path.join(root, f) for f in files]
+        relative_paths = [os.path.relpath(p, dir_path) for p in absolute_paths]
+
+        if add_base_to_relative:
+            relative_paths = [os.path.join(base_name, p) for p in relative_paths]
+
+        assert len(absolute_paths) == len(relative_paths)
+        result += zip(absolute_paths, relative_paths)
+
+    return result
+
+
+def copy_files_and_create_dirs(files: List[Tuple[str, str]]) -> None:
+    """Takes in a list of tuples of (src, dst) paths and copies files.
+    Will create all necessary directories."""
+    for file in files:
+        target_dir_name = os.path.dirname(file[1])
+
+        # will create all intermediate-level directories
+        if not os.path.exists(target_dir_name):
+            os.makedirs(target_dir_name)
+
+        shutil.copyfile(file[0], file[1])
+
+
+# URL helpers
+# ------------------------------------------------------------------------------------------
+
+def is_url(obj: Any, allow_file_urls: bool = False) -> bool:
+    """Determine whether the given object is a valid URL string."""
+    if not isinstance(obj, str) or not "://" in obj:
+        return False
+    if allow_file_urls and obj.startswith('file://'):
+        return True
+    try:
+        res = requests.compat.urlparse(obj)
+        if not res.scheme or not res.netloc or not "." in res.netloc:
+            return False
+        res = requests.compat.urlparse(requests.compat.urljoin(obj, "/"))
+        if not res.scheme or not res.netloc or not "." in res.netloc:
+            return False
+    except:
+        return False
+    return True
+
+
+def open_url(url: str, cache_dir: str = None, num_attempts: int = 10, verbose: bool = True, return_filename: bool = False, cache: bool = True) -> Any:
+    """Download the given URL and return a binary-mode file object to access the data."""
+    assert num_attempts >= 1
+    assert not (return_filename and (not cache))
+
+    # Doesn't look like an URL scheme so interpret it as a local filename.
+    if not re.match('^[a-z]+://', url):
+        return url if return_filename else open(url, "rb")
+
+    # Handle file URLs.  This code handles unusual file:// patterns that
+    # arise on Windows:
+    #
+    # file:///c:/foo.txt
+    #
+    # which would translate to a local '/c:/foo.txt' filename that's
+    # invalid.  Drop the forward slash for such pathnames.
+    #
+    # If you touch this code path, you should test it on both Linux and
+    # Windows.
+    #
+    # Some internet resources suggest using urllib.request.url2pathname() but
+    # but that converts forward slashes to backslashes and this causes
+    # its own set of problems.
+    if url.startswith('file://'):
+        filename = urllib.parse.urlparse(url).path
+        if re.match(r'^/[a-zA-Z]:', filename):
+            filename = filename[1:]
+        return filename if return_filename else open(filename, "rb")
+
+    assert is_url(url)
+
+    # Lookup from cache.
+    if cache_dir is None:
+        cache_dir = make_cache_dir_path('downloads')
+
+    url_md5 = hashlib.md5(url.encode("utf-8")).hexdigest()
+    if cache:
+        cache_files = glob.glob(os.path.join(cache_dir, url_md5 + "_*"))
+        if len(cache_files) == 1:
+            filename = cache_files[0]
+            return filename if return_filename else open(filename, "rb")
+
+    # Download.
+    url_name = None
+    url_data = None
+    with requests.Session() as session:
+        if verbose:
+            print("Downloading %s ..." % url, end="", flush=True)
+        for attempts_left in reversed(range(num_attempts)):
+            try:
+                with session.get(url) as res:
+                    res.raise_for_status()
+                    if len(res.content) == 0:
+                        raise IOError("No data received")
+
+                    if len(res.content) < 8192:
+                        content_str = res.content.decode("utf-8")
+                        if "download_warning" in res.headers.get("Set-Cookie", ""):
+                            links = [html.unescape(link) for link in content_str.split('"') if "export=download" in link]
+                            if len(links) == 1:
+                                url = requests.compat.urljoin(url, links[0])
+                                raise IOError("Google Drive virus checker nag")
+                        if "Google Drive - Quota exceeded" in content_str:
+                            raise IOError("Google Drive download quota exceeded -- please try again later")
+
+                    match = re.search(r'filename="([^"]*)"', res.headers.get("Content-Disposition", ""))
+                    url_name = match[1] if match else url
+                    url_data = res.content
+                    if verbose:
+                        print(" done")
+                    break
+            except KeyboardInterrupt:
+                raise
+            except:
+                if not attempts_left:
+                    if verbose:
+                        print(" failed")
+                    raise
+                if verbose:
+                    print(".", end="", flush=True)
+
+    # Save to cache.
+    if cache:
+        safe_name = re.sub(r"[^0-9a-zA-Z-._]", "_", url_name)
+        cache_file = os.path.join(cache_dir, url_md5 + "_" + safe_name)
+        temp_file = os.path.join(cache_dir, "tmp_" + uuid.uuid4().hex + "_" + url_md5 + "_" + safe_name)
+        os.makedirs(cache_dir, exist_ok=True)
+        with open(temp_file, "wb") as f:
+            f.write(url_data)
+        os.replace(temp_file, cache_file) # atomic
+        if return_filename:
+            return cache_file
+
+    # Return data as file object.
+    assert not return_filename
+    return io.BytesIO(url_data)
+
+
+def dividable(n, k=2):
+    if k == 2:
+        for i in range(int(np.sqrt(n)), 0, -1):
+            if n % i == 0:
+                break
+        return i, n // i
+    elif k == 3:
+        for i in range(int(float(n) ** (1/3)), 0, -1):
+            if n % i == 0:
+                b, c = dividable(n // i, 2)
+                return i, b, c
+    else:
+        raise NotImplementedError
+
+
+def visualize_feature_map(x, scale=1.0, mask=None, loc=None):
+    B, C, H, W = x.size()
+    lh, lw = dividable(C)
+    x = x.reshape(B, lh, lw, H, W).permute(0,1,3,2,4)
+    # loc = [(3,1), (6,3), (4,0)]
+    # loc = [(4,0), (0,7), (6,2)]
+    loc = [(3, 11), (5,3), (3,9)]
+    # loc = [(1,3), (5,3), (7,4)]
+    # loc = [(0,5), (10,0), (3,14)]
+    if loc is None:
+        x = x.reshape(B, 1, lh*H, lw*W).repeat(1,3,1,1)
+    else:
+        x = [x[:, l[0], :, l[1]] for l in loc]
+        x = torch.stack(x, 1)
+        x = x / x.norm(dim=1, keepdim=True)
+    x = x / scale
+    return x
+
+
+def hash_func(x, res, T):
+    d = x.size(-1)
+    assert d <= 3
+
+    h = x[..., 0]
+    if res ** d < T:
+        f = [1, res, res * res]
+        for i in range(1, d):
+            h += x[..., i] * f[i]
+    else:
+        f = [1, 19349663, 83492791]
+        for i in range(1, d):
+            h = torch.bitwise_xor(h, x[..., i] * f[i])
+        h = h % T
+    return h

generate.py

@@ -0,0 +1,202 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+"""Generate images using pretrained network pickle."""
+
+import os
+import re
+import time
+import glob
+from typing import List, Optional
+
+import click
+import dnnlib
+import numpy as np
+import PIL.Image
+import torch
+import imageio
+import legacy
+from renderer import Renderer
+
+#----------------------------------------------------------------------------
+
+def num_range(s: str) -> List[int]:
+    '''Accept either a comma separated list of numbers 'a,b,c' or a range 'a-c' and return as a list of ints.'''
+
+    range_re = re.compile(r'^(\d+)-(\d+)$')
+    m = range_re.match(s)
+    if m:
+        return list(range(int(m.group(1)), int(m.group(2))+1))
+    vals = s.split(',')
+    return [int(x) for x in vals]
+
+#----------------------------------------------------------------------------
+os.environ['PYOPENGL_PLATFORM'] = 'egl'
+
+@click.command()
+@click.pass_context
+@click.option('--network', 'network_pkl', help='Network pickle filename', required=True)
+@click.option('--seeds', type=num_range, help='List of random seeds')
+@click.option('--trunc', 'truncation_psi', type=float, help='Truncation psi', default=1, show_default=True)
+@click.option('--class', 'class_idx', type=int, help='Class label (unconditional if not specified)')
+@click.option('--noise-mode', help='Noise mode', type=click.Choice(['const', 'random', 'none']), default='const', show_default=True)
+@click.option('--projected-w', help='Projection result file', type=str, metavar='FILE')
+@click.option('--outdir', help='Where to save the output images', type=str, required=True, metavar='DIR')
+@click.option('--render-program', default=None, show_default=True)
+@click.option('--render-option', default=None, type=str, help="e.g. up_256, camera, depth")
+@click.option('--n_steps', default=8, type=int, help="number of steps for each seed")
+@click.option('--no-video', default=False)
+@click.option('--relative_range_u_scale', default=1.0, type=float, help="relative scale on top of the original range u")
+def generate_images(
+    ctx: click.Context,
+    network_pkl: str,
+    seeds: Optional[List[int]],
+    truncation_psi: float,
+    noise_mode: str,
+    outdir: str,
+    class_idx: Optional[int],
+    projected_w: Optional[str],
+    render_program=None,
+    render_option=None,
+    n_steps=8,
+    no_video=False,
+    relative_range_u_scale=1.0
+):
+
+    
+    device = torch.device('cuda')
+    if os.path.isdir(network_pkl):
+        network_pkl = sorted(glob.glob(network_pkl + '/*.pkl'))[-1]
+    print('Loading networks from "%s"...' % network_pkl)
+    
+    with dnnlib.util.open_url(network_pkl) as f:
+        network = legacy.load_network_pkl(f)
+        G = network['G_ema'].to(device) # type: ignore
+        D = network['D'].to(device)
+    # from fairseq import pdb;pdb.set_trace()
+    os.makedirs(outdir, exist_ok=True)
+
+    # Labels.
+    label = torch.zeros([1, G.c_dim], device=device)
+    if G.c_dim != 0:
+        if class_idx is None:
+            ctx.fail('Must specify class label with --class when using a conditional network')
+        label[:, class_idx] = 1
+    else:
+        if class_idx is not None:
+            print ('warn: --class=lbl ignored when running on an unconditional network')
+
+    # avoid persistent classes... 
+    from training.networks import Generator
+    # from training.stylenerf import Discriminator
+    from torch_utils import misc
+    with torch.no_grad():
+        G2 = Generator(*G.init_args, **G.init_kwargs).to(device)
+        misc.copy_params_and_buffers(G, G2, require_all=False)
+        # D2 = Discriminator(*D.init_args, **D.init_kwargs).to(device)
+        # misc.copy_params_and_buffers(D, D2, require_all=False)
+    G2 = Renderer(G2, D, program=render_program)
+    
+    # Generate images.
+    all_imgs = []
+
+    def stack_imgs(imgs):
+        img = torch.stack(imgs, dim=2)
+        return img.reshape(img.size(0) * img.size(1), img.size(2) * img.size(3), 3)
+
+    def proc_img(img): 
+        return (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8).cpu()
+
+    if projected_w is not None:
+        ws = np.load(projected_w)
+        ws = torch.tensor(ws, device=device) # pylint: disable=not-callable
+        img = G2(styles=ws, truncation_psi=truncation_psi, noise_mode=noise_mode, render_option=render_option)
+        assert isinstance(img, List)
+        imgs = [proc_img(i) for i in img]
+        all_imgs += [imgs]
+    
+    else:
+        for seed_idx, seed in enumerate(seeds):
+            print('Generating image for seed %d (%d/%d) ...' % (seed, seed_idx, len(seeds)))
+            G2.set_random_seed(seed)
+            z = torch.from_numpy(np.random.RandomState(seed).randn(2, G.z_dim)).to(device)
+            relative_range_u = [0.5 - 0.5 * relative_range_u_scale, 0.5 + 0.5 * relative_range_u_scale]
+            outputs = G2(
+                z=z,
+                c=label,
+                truncation_psi=truncation_psi,
+                noise_mode=noise_mode,
+                render_option=render_option,
+                n_steps=n_steps,
+                relative_range_u=relative_range_u,
+                return_cameras=True)
+            if isinstance(outputs, tuple):
+                img, cameras = outputs
+            else:
+                img = outputs
+
+            if isinstance(img, List):
+                imgs = [proc_img(i) for i in img]
+                if not no_video:
+                    all_imgs += [imgs]
+           
+                curr_out_dir = os.path.join(outdir, 'seed_{:0>6d}'.format(seed))
+                os.makedirs(curr_out_dir, exist_ok=True)
+
+                if (render_option is not None) and ("gen_ibrnet_metadata" in render_option):
+                    intrinsics = []
+                    poses = []
+                    _, H, W, _ = imgs[0].shape
+                    for i, camera in enumerate(cameras):
+                        intri, pose, _, _ = camera
+                        focal = (H - 1) * 0.5 / intri[0, 0, 0].item()
+                        intri = np.diag([focal, focal, 1.0, 1.0]).astype(np.float32)
+                        intri[0, 2], intri[1, 2] = (W - 1) * 0.5, (H - 1) * 0.5
+
+                        pose = pose.squeeze().detach().cpu().numpy() @ np.diag([1, -1, -1, 1]).astype(np.float32)
+                        intrinsics.append(intri)
+                        poses.append(pose)
+
+                    intrinsics = np.stack(intrinsics, axis=0)
+                    poses = np.stack(poses, axis=0)
+
+                    np.savez(os.path.join(curr_out_dir, 'cameras.npz'), intrinsics=intrinsics, poses=poses)
+                    with open(os.path.join(curr_out_dir, 'meta.conf'), 'w') as f:
+                        f.write('depth_range = {}\ntest_hold_out = {}\nheight = {}\nwidth = {}'.
+                                format(G2.generator.synthesis.depth_range, 2, H, W))
+
+                img_dir = os.path.join(curr_out_dir, 'images_raw')
+                os.makedirs(img_dir, exist_ok=True)
+                for step, img in enumerate(imgs):
+                    PIL.Image.fromarray(img[0].detach().cpu().numpy(), 'RGB').save(f'{img_dir}/{step:03d}.png')
+
+            else:
+                img = proc_img(img)[0]
+                PIL.Image.fromarray(img.numpy(), 'RGB').save(f'{outdir}/seed_{seed:0>6d}.png')
+
+    if len(all_imgs) > 0 and (not no_video):
+         # write to video
+        timestamp = time.strftime('%Y%m%d.%H%M%S',time.localtime(time.time()))
+        seeds = ','.join([str(s) for s in seeds]) if seeds is not None else 'projected'
+        network_pkl = network_pkl.split('/')[-1].split('.')[0]
+        all_imgs = [stack_imgs([a[k] for a in all_imgs]).numpy() for k in range(len(all_imgs[0]))]
+        imageio.mimwrite(f'{outdir}/{network_pkl}_{timestamp}_{seeds}.mp4', all_imgs, fps=30, quality=8)
+        outdir = f'{outdir}/{network_pkl}_{timestamp}_{seeds}'
+        os.makedirs(outdir, exist_ok=True)
+        for step, img in enumerate(all_imgs):
+            PIL.Image.fromarray(img, 'RGB').save(f'{outdir}/{step:04d}.png')
+
+
+#----------------------------------------------------------------------------
+
+if __name__ == "__main__":
+    generate_images() # pylint: disable=no-value-for-parameter
+
+#----------------------------------------------------------------------------

gui_utils/__init__.py

@@ -0,0 +1,9 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+# empty

gui_utils/gl_utils.py

@@ -0,0 +1,374 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import os
+import functools
+import contextlib
+import numpy as np
+import OpenGL.GL as gl
+import OpenGL.GL.ARB.texture_float
+import dnnlib
+
+#----------------------------------------------------------------------------
+
+def init_egl():
+    assert os.environ['PYOPENGL_PLATFORM'] == 'egl' # Must be set before importing OpenGL.
+    import OpenGL.EGL as egl
+    import ctypes
+
+    # Initialize EGL.
+    display = egl.eglGetDisplay(egl.EGL_DEFAULT_DISPLAY)
+    assert display != egl.EGL_NO_DISPLAY
+    major = ctypes.c_int32()
+    minor = ctypes.c_int32()
+    ok = egl.eglInitialize(display, major, minor)
+    assert ok
+    assert major.value * 10 + minor.value >= 14
+
+    # Choose config.
+    config_attribs = [
+        egl.EGL_RENDERABLE_TYPE,    egl.EGL_OPENGL_BIT,
+        egl.EGL_SURFACE_TYPE,       egl.EGL_PBUFFER_BIT,
+        egl.EGL_NONE
+    ]
+    configs = (ctypes.c_int32 * 1)()
+    num_configs = ctypes.c_int32()
+    ok = egl.eglChooseConfig(display, config_attribs, configs, 1, num_configs)
+    assert ok
+    assert num_configs.value == 1
+    config = configs[0]
+
+    # Create dummy pbuffer surface.
+    surface_attribs = [
+        egl.EGL_WIDTH,  1,
+        egl.EGL_HEIGHT, 1,
+        egl.EGL_NONE
+    ]
+    surface = egl.eglCreatePbufferSurface(display, config, surface_attribs)
+    assert surface != egl.EGL_NO_SURFACE
+
+    # Setup GL context.
+    ok = egl.eglBindAPI(egl.EGL_OPENGL_API)
+    assert ok
+    context = egl.eglCreateContext(display, config, egl.EGL_NO_CONTEXT, None)
+    assert context != egl.EGL_NO_CONTEXT
+    ok = egl.eglMakeCurrent(display, surface, surface, context)
+    assert ok
+
+#----------------------------------------------------------------------------
+
+_texture_formats = {
+    ('uint8',   1): dnnlib.EasyDict(type=gl.GL_UNSIGNED_BYTE, format=gl.GL_LUMINANCE,       internalformat=gl.GL_LUMINANCE8),
+    ('uint8',   2): dnnlib.EasyDict(type=gl.GL_UNSIGNED_BYTE, format=gl.GL_LUMINANCE_ALPHA, internalformat=gl.GL_LUMINANCE8_ALPHA8),
+    ('uint8',   3): dnnlib.EasyDict(type=gl.GL_UNSIGNED_BYTE, format=gl.GL_RGB,             internalformat=gl.GL_RGB8),
+    ('uint8',   4): dnnlib.EasyDict(type=gl.GL_UNSIGNED_BYTE, format=gl.GL_RGBA,            internalformat=gl.GL_RGBA8),
+    ('float32', 1): dnnlib.EasyDict(type=gl.GL_FLOAT,         format=gl.GL_LUMINANCE,       internalformat=OpenGL.GL.ARB.texture_float.GL_LUMINANCE32F_ARB),
+    ('float32', 2): dnnlib.EasyDict(type=gl.GL_FLOAT,         format=gl.GL_LUMINANCE_ALPHA, internalformat=OpenGL.GL.ARB.texture_float.GL_LUMINANCE_ALPHA32F_ARB),
+    ('float32', 3): dnnlib.EasyDict(type=gl.GL_FLOAT,         format=gl.GL_RGB,             internalformat=gl.GL_RGB32F),
+    ('float32', 4): dnnlib.EasyDict(type=gl.GL_FLOAT,         format=gl.GL_RGBA,            internalformat=gl.GL_RGBA32F),
+}
+
+def get_texture_format(dtype, channels):
+    return _texture_formats[(np.dtype(dtype).name, int(channels))]
+
+#----------------------------------------------------------------------------
+
+def prepare_texture_data(image):
+    image = np.asarray(image)
+    if image.ndim == 2:
+        image = image[:, :, np.newaxis]
+    if image.dtype.name == 'float64':
+        image = image.astype('float32')
+    return image
+
+#----------------------------------------------------------------------------
+
+def draw_pixels(image, *, pos=0, zoom=1, align=0, rint=True):
+    pos = np.broadcast_to(np.asarray(pos, dtype='float32'), [2])
+    zoom = np.broadcast_to(np.asarray(zoom, dtype='float32'), [2])
+    align = np.broadcast_to(np.asarray(align, dtype='float32'), [2])
+    image = prepare_texture_data(image)
+    height, width, channels = image.shape
+    size = zoom * [width, height]
+    pos = pos - size * align
+    if rint:
+        pos = np.rint(pos)
+    fmt = get_texture_format(image.dtype, channels)
+
+    gl.glPushAttrib(gl.GL_CURRENT_BIT | gl.GL_PIXEL_MODE_BIT)
+    gl.glPushClientAttrib(gl.GL_CLIENT_PIXEL_STORE_BIT)
+    gl.glRasterPos2f(pos[0], pos[1])
+    gl.glPixelZoom(zoom[0], -zoom[1])
+    gl.glPixelStorei(gl.GL_UNPACK_ALIGNMENT, 1)
+    gl.glDrawPixels(width, height, fmt.format, fmt.type, image)
+    gl.glPopClientAttrib()
+    gl.glPopAttrib()
+
+#----------------------------------------------------------------------------
+
+def read_pixels(width, height, *, pos=0, dtype='uint8', channels=3):
+    pos = np.broadcast_to(np.asarray(pos, dtype='float32'), [2])
+    dtype = np.dtype(dtype)
+    fmt = get_texture_format(dtype, channels)
+    image = np.empty([height, width, channels], dtype=dtype)
+
+    gl.glPushClientAttrib(gl.GL_CLIENT_PIXEL_STORE_BIT)
+    gl.glPixelStorei(gl.GL_PACK_ALIGNMENT, 1)
+    gl.glReadPixels(int(np.round(pos[0])), int(np.round(pos[1])), width, height, fmt.format, fmt.type, image)
+    gl.glPopClientAttrib()
+    return np.flipud(image)
+
+#----------------------------------------------------------------------------
+
+class Texture:
+    def __init__(self, *, image=None, width=None, height=None, channels=None, dtype=None, bilinear=True, mipmap=True):
+        self.gl_id = None
+        self.bilinear = bilinear
+        self.mipmap = mipmap
+
+        # Determine size and dtype.
+        if image is not None:
+            image = prepare_texture_data(image)
+            self.height, self.width, self.channels = image.shape
+            self.dtype = image.dtype
+        else:
+            assert width is not None and height is not None
+            self.width = width
+            self.height = height
+            self.channels = channels if channels is not None else 3
+            self.dtype = np.dtype(dtype) if dtype is not None else np.uint8
+
+        # Validate size and dtype.
+        assert isinstance(self.width, int) and self.width >= 0
+        assert isinstance(self.height, int) and self.height >= 0
+        assert isinstance(self.channels, int) and self.channels >= 1
+        assert self.is_compatible(width=width, height=height, channels=channels, dtype=dtype)
+
+        # Create texture object.
+        self.gl_id = gl.glGenTextures(1)
+        with self.bind():
+            gl.glTexParameterf(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_WRAP_S, gl.GL_CLAMP_TO_EDGE)
+            gl.glTexParameterf(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_WRAP_T, gl.GL_CLAMP_TO_EDGE)
+            gl.glTexParameterf(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_MAG_FILTER, gl.GL_LINEAR if self.bilinear else gl.GL_NEAREST)
+            gl.glTexParameterf(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_MIN_FILTER, gl.GL_LINEAR_MIPMAP_LINEAR if self.mipmap else gl.GL_NEAREST)
+        self.update(image)
+
+    def delete(self):
+        if self.gl_id is not None:
+            gl.glDeleteTextures([self.gl_id])
+            self.gl_id = None
+
+    def __del__(self):
+        try:
+            self.delete()
+        except:
+            pass
+
+    @contextlib.contextmanager
+    def bind(self):
+        prev_id = gl.glGetInteger(gl.GL_TEXTURE_BINDING_2D)
+        gl.glBindTexture(gl.GL_TEXTURE_2D, self.gl_id)
+        yield
+        gl.glBindTexture(gl.GL_TEXTURE_2D, prev_id)
+
+    def update(self, image):
+        if image is not None:
+            image = prepare_texture_data(image)
+            assert self.is_compatible(image=image)
+        with self.bind():
+            fmt = get_texture_format(self.dtype, self.channels)
+            gl.glPushClientAttrib(gl.GL_CLIENT_PIXEL_STORE_BIT)
+            gl.glPixelStorei(gl.GL_UNPACK_ALIGNMENT, 1)
+            gl.glTexImage2D(gl.GL_TEXTURE_2D, 0, fmt.internalformat, self.width, self.height, 0, fmt.format, fmt.type, image)
+            if self.mipmap:
+                gl.glGenerateMipmap(gl.GL_TEXTURE_2D)
+            gl.glPopClientAttrib()
+
+    def draw(self, *, pos=0, zoom=1, align=0, rint=False, color=1, alpha=1, rounding=0):
+        zoom = np.broadcast_to(np.asarray(zoom, dtype='float32'), [2])
+        size = zoom * [self.width, self.height]
+        with self.bind():
+            gl.glPushAttrib(gl.GL_ENABLE_BIT)
+            gl.glEnable(gl.GL_TEXTURE_2D)
+            draw_rect(pos=pos, size=size, align=align, rint=rint, color=color, alpha=alpha, rounding=rounding)
+            gl.glPopAttrib()
+
+    def is_compatible(self, *, image=None, width=None, height=None, channels=None, dtype=None): # pylint: disable=too-many-return-statements
+        if image is not None:
+            if image.ndim != 3:
+                return False
+            ih, iw, ic = image.shape
+            if not self.is_compatible(width=iw, height=ih, channels=ic, dtype=image.dtype):
+                return False
+        if width is not None and self.width != width:
+            return False
+        if height is not None and self.height != height:
+            return False
+        if channels is not None and self.channels != channels:
+            return False
+        if dtype is not None and self.dtype != dtype:
+            return False
+        return True
+
+#----------------------------------------------------------------------------
+
+class Framebuffer:
+    def __init__(self, *, texture=None, width=None, height=None, channels=None, dtype=None, msaa=0):
+        self.texture = texture
+        self.gl_id = None
+        self.gl_color = None
+        self.gl_depth_stencil = None
+        self.msaa = msaa
+
+        # Determine size and dtype.
+        if texture is not None:
+            assert isinstance(self.texture, Texture)
+            self.width = texture.width
+            self.height = texture.height
+            self.channels = texture.channels
+            self.dtype = texture.dtype
+        else:
+            assert width is not None and height is not None
+            self.width = width
+            self.height = height
+            self.channels = channels if channels is not None else 4
+            self.dtype = np.dtype(dtype) if dtype is not None else np.float32
+
+        # Validate size and dtype.
+        assert isinstance(self.width, int) and self.width >= 0
+        assert isinstance(self.height, int) and self.height >= 0
+        assert isinstance(self.channels, int) and self.channels >= 1
+        assert width is None or width == self.width
+        assert height is None or height == self.height
+        assert channels is None or channels == self.channels
+        assert dtype is None or dtype == self.dtype
+
+        # Create framebuffer object.
+        self.gl_id = gl.glGenFramebuffers(1)
+        with self.bind():
+
+            # Setup color buffer.
+            if self.texture is not None:
+                assert self.msaa == 0
+                gl.glFramebufferTexture2D(gl.GL_FRAMEBUFFER, gl.GL_COLOR_ATTACHMENT0, gl.GL_TEXTURE_2D, self.texture.gl_id, 0)
+            else:
+                fmt = get_texture_format(self.dtype, self.channels)
+                self.gl_color = gl.glGenRenderbuffers(1)
+                gl.glBindRenderbuffer(gl.GL_RENDERBUFFER, self.gl_color)
+                gl.glRenderbufferStorageMultisample(gl.GL_RENDERBUFFER, self.msaa, fmt.internalformat, self.width, self.height)
+                gl.glFramebufferRenderbuffer(gl.GL_FRAMEBUFFER, gl.GL_COLOR_ATTACHMENT0, gl.GL_RENDERBUFFER, self.gl_color)
+
+            # Setup depth/stencil buffer.
+            self.gl_depth_stencil = gl.glGenRenderbuffers(1)
+            gl.glBindRenderbuffer(gl.GL_RENDERBUFFER, self.gl_depth_stencil)
+            gl.glRenderbufferStorageMultisample(gl.GL_RENDERBUFFER, self.msaa, gl.GL_DEPTH24_STENCIL8, self.width, self.height)
+            gl.glFramebufferRenderbuffer(gl.GL_FRAMEBUFFER, gl.GL_DEPTH_STENCIL_ATTACHMENT, gl.GL_RENDERBUFFER, self.gl_depth_stencil)
+
+    def delete(self):
+        if self.gl_id is not None:
+            gl.glDeleteFramebuffers([self.gl_id])
+            self.gl_id = None
+        if self.gl_color is not None:
+            gl.glDeleteRenderbuffers(1, [self.gl_color])
+            self.gl_color = None
+        if self.gl_depth_stencil is not None:
+            gl.glDeleteRenderbuffers(1, [self.gl_depth_stencil])
+            self.gl_depth_stencil = None
+
+    def __del__(self):
+        try:
+            self.delete()
+        except:
+            pass
+
+    @contextlib.contextmanager
+    def bind(self):
+        prev_fbo = gl.glGetInteger(gl.GL_FRAMEBUFFER_BINDING)
+        prev_rbo = gl.glGetInteger(gl.GL_RENDERBUFFER_BINDING)
+        gl.glBindFramebuffer(gl.GL_FRAMEBUFFER, self.gl_id)
+        if self.width is not None and self.height is not None:
+            gl.glViewport(0, 0, self.width, self.height)
+        yield
+        gl.glBindFramebuffer(gl.GL_FRAMEBUFFER, prev_fbo)
+        gl.glBindRenderbuffer(gl.GL_RENDERBUFFER, prev_rbo)
+
+    def blit(self, dst=None):
+        assert dst is None or isinstance(dst, Framebuffer)
+        with self.bind():
+            gl.glBindFramebuffer(gl.GL_DRAW_FRAMEBUFFER, 0 if dst is None else dst.fbo)
+            gl.glBlitFramebuffer(0, 0, self.width, self.height, 0, 0, self.width, self.height, gl.GL_COLOR_BUFFER_BIT, gl.GL_NEAREST)
+
+#----------------------------------------------------------------------------
+
+def draw_shape(vertices, *, mode=gl.GL_TRIANGLE_FAN, pos=0, size=1, color=1, alpha=1):
+    assert vertices.ndim == 2 and vertices.shape[1] == 2
+    pos = np.broadcast_to(np.asarray(pos, dtype='float32'), [2])
+    size = np.broadcast_to(np.asarray(size, dtype='float32'), [2])
+    color = np.broadcast_to(np.asarray(color, dtype='float32'), [3])
+    alpha = np.clip(np.broadcast_to(np.asarray(alpha, dtype='float32'), []), 0, 1)
+
+    gl.glPushClientAttrib(gl.GL_CLIENT_VERTEX_ARRAY_BIT)
+    gl.glPushAttrib(gl.GL_CURRENT_BIT | gl.GL_TRANSFORM_BIT)
+    gl.glMatrixMode(gl.GL_MODELVIEW)
+    gl.glPushMatrix()
+
+    gl.glEnableClientState(gl.GL_VERTEX_ARRAY)
+    gl.glEnableClientState(gl.GL_TEXTURE_COORD_ARRAY)
+    gl.glVertexPointer(2, gl.GL_FLOAT, 0, vertices)
+    gl.glTexCoordPointer(2, gl.GL_FLOAT, 0, vertices)
+    gl.glTranslate(pos[0], pos[1], 0)
+    gl.glScale(size[0], size[1], 1)
+    gl.glColor4f(color[0] * alpha, color[1] * alpha, color[2] * alpha, alpha)
+    gl.glDrawArrays(mode, 0, vertices.shape[0])
+
+    gl.glPopMatrix()
+    gl.glPopAttrib()
+    gl.glPopClientAttrib()
+
+#----------------------------------------------------------------------------
+
+def draw_rect(*, pos=0, pos2=None, size=None, align=0, rint=False, color=1, alpha=1, rounding=0):
+    assert pos2 is None or size is None
+    pos = np.broadcast_to(np.asarray(pos, dtype='float32'), [2])
+    pos2 = np.broadcast_to(np.asarray(pos2, dtype='float32'), [2]) if pos2 is not None else None
+    size = np.broadcast_to(np.asarray(size, dtype='float32'), [2]) if size is not None else None
+    size = size if size is not None else pos2 - pos if pos2 is not None else np.array([1, 1], dtype='float32')
+    pos = pos - size * align
+    if rint:
+        pos = np.rint(pos)
+    rounding = np.broadcast_to(np.asarray(rounding, dtype='float32'), [2])
+    rounding = np.minimum(np.abs(rounding) / np.maximum(np.abs(size), 1e-8), 0.5)
+    if np.min(rounding) == 0:
+        rounding *= 0
+    vertices = _setup_rect(float(rounding[0]), float(rounding[1]))
+    draw_shape(vertices, mode=gl.GL_TRIANGLE_FAN, pos=pos, size=size, color=color, alpha=alpha)
+
+@functools.lru_cache(maxsize=10000)
+def _setup_rect(rx, ry):
+    t = np.linspace(0, np.pi / 2, 1 if max(rx, ry) == 0 else 64)
+    s = 1 - np.sin(t); c = 1 - np.cos(t)
+    x = [c * rx, 1 - s * rx, 1 - c * rx, s * rx]
+    y = [s * ry, c * ry, 1 - s * ry, 1 - c * ry]
+    v = np.stack([x, y], axis=-1).reshape(-1, 2)
+    return v.astype('float32')
+
+#----------------------------------------------------------------------------
+
+def draw_circle(*, center=0, radius=100, hole=0, color=1, alpha=1):
+    hole = np.broadcast_to(np.asarray(hole, dtype='float32'), [])
+    vertices = _setup_circle(float(hole))
+    draw_shape(vertices, mode=gl.GL_TRIANGLE_STRIP, pos=center, size=radius, color=color, alpha=alpha)
+
+@functools.lru_cache(maxsize=10000)
+def _setup_circle(hole):
+    t = np.linspace(0, np.pi * 2, 128)
+    s = np.sin(t); c = np.cos(t)
+    v = np.stack([c, s, c * hole, s * hole], axis=-1).reshape(-1, 2)
+    return v.astype('float32')
+
+#----------------------------------------------------------------------------

gui_utils/glfw_window.py

@@ -0,0 +1,229 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import time
+import glfw
+import OpenGL.GL as gl
+from . import gl_utils
+
+#----------------------------------------------------------------------------
+
+class GlfwWindow: # pylint: disable=too-many-public-methods
+    def __init__(self, *, title='GlfwWindow', window_width=1920, window_height=1080, deferred_show=True, close_on_esc=True):
+        self._glfw_window           = None
+        self._drawing_frame         = False
+        self._frame_start_time      = None
+        self._frame_delta           = 0
+        self._fps_limit             = None
+        self._vsync                 = None
+        self._skip_frames           = 0
+        self._deferred_show         = deferred_show
+        self._close_on_esc          = close_on_esc
+        self._esc_pressed           = False
+        self._drag_and_drop_paths   = None
+        self._capture_next_frame    = False
+        self._captured_frame        = None
+
+        # Create window.
+        glfw.init()
+        glfw.window_hint(glfw.VISIBLE, False)
+        self._glfw_window = glfw.create_window(width=window_width, height=window_height, title=title, monitor=None, share=None)
+        self._attach_glfw_callbacks()
+        self.make_context_current()
+
+        # Adjust window.
+        self.set_vsync(False)
+        self.set_window_size(window_width, window_height)
+        if not self._deferred_show:
+            glfw.show_window(self._glfw_window)
+
+    def close(self):
+        if self._drawing_frame:
+            self.end_frame()
+        if self._glfw_window is not None:
+            glfw.destroy_window(self._glfw_window)
+            self._glfw_window = None
+        #glfw.terminate() # Commented out to play it nice with other glfw clients.
+
+    def __del__(self):
+        try:
+            self.close()
+        except:
+            pass
+
+    @property
+    def window_width(self):
+        return self.content_width
+
+    @property
+    def window_height(self):
+        return self.content_height + self.title_bar_height
+
+    @property
+    def content_width(self):
+        width, _height = glfw.get_window_size(self._glfw_window)
+        return width
+
+    @property
+    def content_height(self):
+        _width, height = glfw.get_window_size(self._glfw_window)
+        return height
+
+    @property
+    def title_bar_height(self):
+        _left, top, _right, _bottom = glfw.get_window_frame_size(self._glfw_window)
+        return top
+
+    @property
+    def monitor_width(self):
+        _, _, width, _height = glfw.get_monitor_workarea(glfw.get_primary_monitor())
+        return width
+
+    @property
+    def monitor_height(self):
+        _, _, _width, height = glfw.get_monitor_workarea(glfw.get_primary_monitor())
+        return height
+
+    @property
+    def frame_delta(self):
+        return self._frame_delta
+
+    def set_title(self, title):
+        glfw.set_window_title(self._glfw_window, title)
+
+    def set_window_size(self, width, height):
+        width = min(width, self.monitor_width)
+        height = min(height, self.monitor_height)
+        glfw.set_window_size(self._glfw_window, width, max(height - self.title_bar_height, 0))
+        if width == self.monitor_width and height == self.monitor_height:
+            self.maximize()
+
+    def set_content_size(self, width, height):
+        self.set_window_size(width, height + self.title_bar_height)
+
+    def maximize(self):
+        glfw.maximize_window(self._glfw_window)
+
+    def set_position(self, x, y):
+        glfw.set_window_pos(self._glfw_window, x, y + self.title_bar_height)
+
+    def center(self):
+        self.set_position((self.monitor_width - self.window_width) // 2, (self.monitor_height - self.window_height) // 2)
+
+    def set_vsync(self, vsync):
+        vsync = bool(vsync)
+        if vsync != self._vsync:
+            glfw.swap_interval(1 if vsync else 0)
+            self._vsync = vsync
+
+    def set_fps_limit(self, fps_limit):
+        self._fps_limit = int(fps_limit)
+
+    def should_close(self):
+        return glfw.window_should_close(self._glfw_window) or (self._close_on_esc and self._esc_pressed)
+
+    def skip_frame(self):
+        self.skip_frames(1)
+
+    def skip_frames(self, num): # Do not update window for the next N frames.
+        self._skip_frames = max(self._skip_frames, int(num))
+
+    def is_skipping_frames(self):
+        return self._skip_frames > 0
+
+    def capture_next_frame(self):
+        self._capture_next_frame = True
+
+    def pop_captured_frame(self):
+        frame = self._captured_frame
+        self._captured_frame = None
+        return frame
+
+    def pop_drag_and_drop_paths(self):
+        paths = self._drag_and_drop_paths
+        self._drag_and_drop_paths = None
+        return paths
+
+    def draw_frame(self): # To be overridden by subclass.
+        self.begin_frame()
+        # Rendering code goes here.
+        self.end_frame()
+
+    def make_context_current(self):
+        if self._glfw_window is not None:
+            glfw.make_context_current(self._glfw_window)
+
+    def begin_frame(self):
+        # End previous frame.
+        if self._drawing_frame:
+            self.end_frame()
+
+        # Apply FPS limit.
+        if self._frame_start_time is not None and self._fps_limit is not None:
+            delay = self._frame_start_time - time.perf_counter() + 1 / self._fps_limit
+            if delay > 0:
+                time.sleep(delay)
+        cur_time = time.perf_counter()
+        if self._frame_start_time is not None:
+            self._frame_delta = cur_time - self._frame_start_time
+        self._frame_start_time = cur_time
+
+        # Process events.
+        glfw.poll_events()
+
+        # Begin frame.
+        self._drawing_frame = True
+        self.make_context_current()
+
+        # Initialize GL state.
+        gl.glViewport(0, 0, self.content_width, self.content_height)
+        gl.glMatrixMode(gl.GL_PROJECTION)
+        gl.glLoadIdentity()
+        gl.glTranslate(-1, 1, 0)
+        gl.glScale(2 / max(self.content_width, 1), -2 / max(self.content_height, 1), 1)
+        gl.glMatrixMode(gl.GL_MODELVIEW)
+        gl.glLoadIdentity()
+        gl.glEnable(gl.GL_BLEND)
+        gl.glBlendFunc(gl.GL_ONE, gl.GL_ONE_MINUS_SRC_ALPHA) # Pre-multiplied alpha.
+
+        # Clear.
+        gl.glClearColor(0, 0, 0, 1)
+        gl.glClear(gl.GL_COLOR_BUFFER_BIT | gl.GL_DEPTH_BUFFER_BIT)
+
+    def end_frame(self):
+        assert self._drawing_frame
+        self._drawing_frame = False
+
+        # Skip frames if requested.
+        if self._skip_frames > 0:
+            self._skip_frames -= 1
+            return
+
+        # Capture frame if requested.
+        if self._capture_next_frame:
+            self._captured_frame = gl_utils.read_pixels(self.content_width, self.content_height)
+            self._capture_next_frame = False
+
+        # Update window.
+        if self._deferred_show:
+            glfw.show_window(self._glfw_window)
+            self._deferred_show = False
+        glfw.swap_buffers(self._glfw_window)
+
+    def _attach_glfw_callbacks(self):
+        glfw.set_key_callback(self._glfw_window, self._glfw_key_callback)
+        glfw.set_drop_callback(self._glfw_window, self._glfw_drop_callback)
+
+    def _glfw_key_callback(self, _window, key, _scancode, action, _mods):
+        if action == glfw.PRESS and key == glfw.KEY_ESCAPE:
+            self._esc_pressed = True
+
+    def _glfw_drop_callback(self, _window, paths):
+        self._drag_and_drop_paths = paths
+
+#----------------------------------------------------------------------------

gui_utils/imgui_utils.py

@@ -0,0 +1,169 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import contextlib
+import imgui
+
+#----------------------------------------------------------------------------
+
+def set_default_style(color_scheme='dark', spacing=9, indent=23, scrollbar=27):
+    s = imgui.get_style()
+    s.window_padding        = [spacing, spacing]
+    s.item_spacing          = [spacing, spacing]
+    s.item_inner_spacing    = [spacing, spacing]
+    s.columns_min_spacing   = spacing
+    s.indent_spacing        = indent
+    s.scrollbar_size        = scrollbar
+    s.frame_padding         = [4, 3]
+    s.window_border_size    = 1
+    s.child_border_size     = 1
+    s.popup_border_size     = 1
+    s.frame_border_size     = 1
+    s.window_rounding       = 0
+    s.child_rounding        = 0
+    s.popup_rounding        = 3
+    s.frame_rounding        = 3
+    s.scrollbar_rounding    = 3
+    s.grab_rounding         = 3
+
+    getattr(imgui, f'style_colors_{color_scheme}')(s)
+    c0 = s.colors[imgui.COLOR_MENUBAR_BACKGROUND]
+    c1 = s.colors[imgui.COLOR_FRAME_BACKGROUND]
+    s.colors[imgui.COLOR_POPUP_BACKGROUND] = [x * 0.7 + y * 0.3 for x, y in zip(c0, c1)][:3] + [1]
+
+#----------------------------------------------------------------------------
+
+@contextlib.contextmanager
+def grayed_out(cond=True):
+    if cond:
+        s = imgui.get_style()
+        text = s.colors[imgui.COLOR_TEXT_DISABLED]
+        grab = s.colors[imgui.COLOR_SCROLLBAR_GRAB]
+        back = s.colors[imgui.COLOR_MENUBAR_BACKGROUND]
+        imgui.push_style_color(imgui.COLOR_TEXT, *text)
+        imgui.push_style_color(imgui.COLOR_CHECK_MARK, *grab)
+        imgui.push_style_color(imgui.COLOR_SLIDER_GRAB, *grab)
+        imgui.push_style_color(imgui.COLOR_SLIDER_GRAB_ACTIVE, *grab)
+        imgui.push_style_color(imgui.COLOR_FRAME_BACKGROUND, *back)
+        imgui.push_style_color(imgui.COLOR_FRAME_BACKGROUND_HOVERED, *back)
+        imgui.push_style_color(imgui.COLOR_FRAME_BACKGROUND_ACTIVE, *back)
+        imgui.push_style_color(imgui.COLOR_BUTTON, *back)
+        imgui.push_style_color(imgui.COLOR_BUTTON_HOVERED, *back)
+        imgui.push_style_color(imgui.COLOR_BUTTON_ACTIVE, *back)
+        imgui.push_style_color(imgui.COLOR_HEADER, *back)
+        imgui.push_style_color(imgui.COLOR_HEADER_HOVERED, *back)
+        imgui.push_style_color(imgui.COLOR_HEADER_ACTIVE, *back)
+        imgui.push_style_color(imgui.COLOR_POPUP_BACKGROUND, *back)
+        yield
+        imgui.pop_style_color(14)
+    else:
+        yield
+
+#----------------------------------------------------------------------------
+
+@contextlib.contextmanager
+def item_width(width=None):
+    if width is not None:
+        imgui.push_item_width(width)
+        yield
+        imgui.pop_item_width()
+    else:
+        yield
+
+#----------------------------------------------------------------------------
+
+def scoped_by_object_id(method):
+    def decorator(self, *args, **kwargs):
+        imgui.push_id(str(id(self)))
+        res = method(self, *args, **kwargs)
+        imgui.pop_id()
+        return res
+    return decorator
+
+#----------------------------------------------------------------------------
+
+def button(label, width=0, enabled=True):
+    with grayed_out(not enabled):
+        clicked = imgui.button(label, width=width)
+    clicked = clicked and enabled
+    return clicked
+
+#----------------------------------------------------------------------------
+
+def collapsing_header(text, visible=None, flags=0, default=False, enabled=True, show=True):
+    expanded = False
+    if show:
+        if default:
+            flags |= imgui.TREE_NODE_DEFAULT_OPEN
+        if not enabled:
+            flags |= imgui.TREE_NODE_LEAF
+        with grayed_out(not enabled):
+            expanded, visible = imgui.collapsing_header(text, visible=visible, flags=flags)
+        expanded = expanded and enabled
+    return expanded, visible
+
+#----------------------------------------------------------------------------
+
+def popup_button(label, width=0, enabled=True):
+    if button(label, width, enabled):
+        imgui.open_popup(label)
+    opened = imgui.begin_popup(label)
+    return opened
+
+#----------------------------------------------------------------------------
+
+def input_text(label, value, buffer_length, flags, width=None, help_text=''):
+    old_value = value
+    color = list(imgui.get_style().colors[imgui.COLOR_TEXT])
+    if value == '':
+        color[-1] *= 0.5
+    with item_width(width):
+        imgui.push_style_color(imgui.COLOR_TEXT, *color)
+        value = value if value != '' else help_text
+        changed, value = imgui.input_text(label, value, buffer_length, flags)
+        value = value if value != help_text else ''
+        imgui.pop_style_color(1)
+    if not flags & imgui.INPUT_TEXT_ENTER_RETURNS_TRUE:
+        changed = (value != old_value)
+    return changed, value
+
+#----------------------------------------------------------------------------
+
+def drag_previous_control(enabled=True):
+    dragging = False
+    dx = 0
+    dy = 0
+    if imgui.begin_drag_drop_source(imgui.DRAG_DROP_SOURCE_NO_PREVIEW_TOOLTIP):
+        if enabled:
+            dragging = True
+            dx, dy = imgui.get_mouse_drag_delta()
+            imgui.reset_mouse_drag_delta()
+        imgui.end_drag_drop_source()
+    return dragging, dx, dy
+
+#----------------------------------------------------------------------------
+
+def drag_button(label, width=0, enabled=True):
+    clicked = button(label, width=width, enabled=enabled)
+    dragging, dx, dy = drag_previous_control(enabled=enabled)
+    return clicked, dragging, dx, dy
+
+#----------------------------------------------------------------------------
+
+def drag_hidden_window(label, x, y, width, height, enabled=True):
+    imgui.push_style_color(imgui.COLOR_WINDOW_BACKGROUND, 0, 0, 0, 0)
+    imgui.push_style_color(imgui.COLOR_BORDER, 0, 0, 0, 0)
+    imgui.set_next_window_position(x, y)
+    imgui.set_next_window_size(width, height)
+    imgui.begin(label, closable=False, flags=(imgui.WINDOW_NO_TITLE_BAR | imgui.WINDOW_NO_RESIZE | imgui.WINDOW_NO_MOVE))
+    dragging, dx, dy = drag_previous_control(enabled=enabled)
+    imgui.end()
+    imgui.pop_style_color(2)
+    return dragging, dx, dy
+
+#----------------------------------------------------------------------------

gui_utils/imgui_window.py

@@ -0,0 +1,103 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import os
+import imgui
+import imgui.integrations.glfw
+
+from . import glfw_window
+from . import imgui_utils
+from . import text_utils
+
+#----------------------------------------------------------------------------
+
+class ImguiWindow(glfw_window.GlfwWindow):
+    def __init__(self, *, title='ImguiWindow', font=None, font_sizes=range(14,24), **glfw_kwargs):
+        if font is None:
+            font = text_utils.get_default_font()
+        font_sizes = {int(size) for size in font_sizes}
+        super().__init__(title=title, **glfw_kwargs)
+
+        # Init fields.
+        self._imgui_context  = None
+        self._imgui_renderer = None
+        self._imgui_fonts    = None
+        self._cur_font_size  = max(font_sizes)
+
+        # Delete leftover imgui.ini to avoid unexpected behavior.
+        if os.path.isfile('imgui.ini'):
+            os.remove('imgui.ini')
+
+        # Init ImGui.
+        self._imgui_context = imgui.create_context()
+        self._imgui_renderer = _GlfwRenderer(self._glfw_window)
+        self._attach_glfw_callbacks()
+        imgui.get_io().ini_saving_rate = 0 # Disable creating imgui.ini at runtime.
+        imgui.get_io().mouse_drag_threshold = 0 # Improve behavior with imgui_utils.drag_custom().
+        self._imgui_fonts = {size: imgui.get_io().fonts.add_font_from_file_ttf(font, size) for size in font_sizes}
+        self._imgui_renderer.refresh_font_texture()
+
+    def close(self):
+        self.make_context_current()
+        self._imgui_fonts = None
+        if self._imgui_renderer is not None:
+            self._imgui_renderer.shutdown()
+            self._imgui_renderer = None
+        if self._imgui_context is not None:
+            #imgui.destroy_context(self._imgui_context) # Commented out to avoid creating imgui.ini at the end.
+            self._imgui_context = None
+        super().close()
+
+    def _glfw_key_callback(self, *args):
+        super()._glfw_key_callback(*args)
+        self._imgui_renderer.keyboard_callback(*args)
+
+    @property
+    def font_size(self):
+        return self._cur_font_size
+
+    @property
+    def spacing(self):
+        return round(self._cur_font_size * 0.4)
+
+    def set_font_size(self, target): # Applied on next frame.
+        self._cur_font_size = min((abs(key - target), key) for key in self._imgui_fonts.keys())[1]
+
+    def begin_frame(self):
+        # Begin glfw frame.
+        super().begin_frame()
+
+        # Process imgui events.
+        self._imgui_renderer.mouse_wheel_multiplier = self._cur_font_size / 10
+        if self.content_width > 0 and self.content_height > 0:
+            self._imgui_renderer.process_inputs()
+
+        # Begin imgui frame.
+        imgui.new_frame()
+        imgui.push_font(self._imgui_fonts[self._cur_font_size])
+        imgui_utils.set_default_style(spacing=self.spacing, indent=self.font_size, scrollbar=self.font_size+4)
+
+    def end_frame(self):
+        imgui.pop_font()
+        imgui.render()
+        imgui.end_frame()
+        self._imgui_renderer.render(imgui.get_draw_data())
+        super().end_frame()
+
+#----------------------------------------------------------------------------
+# Wrapper class for GlfwRenderer to fix a mouse wheel bug on Linux.
+
+class _GlfwRenderer(imgui.integrations.glfw.GlfwRenderer):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.mouse_wheel_multiplier = 1
+
+    def scroll_callback(self, window, x_offset, y_offset):
+        self.io.mouse_wheel += y_offset * self.mouse_wheel_multiplier
+
+#----------------------------------------------------------------------------

gui_utils/text_utils.py

@@ -0,0 +1,123 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import functools
+from typing import Optional
+
+import dnnlib
+import numpy as np
+import PIL.Image
+import PIL.ImageFont
+import scipy.ndimage
+
+from . import gl_utils
+
+#----------------------------------------------------------------------------
+
+def get_default_font():
+    url = 'http://fonts.gstatic.com/s/opensans/v17/mem8YaGs126MiZpBA-U1UpcaXcl0Aw.ttf' # Open Sans regular
+    return dnnlib.util.open_url(url, return_filename=True)
+
+#----------------------------------------------------------------------------
+
+@functools.lru_cache(maxsize=None)
+def get_pil_font(font=None, size=32):
+    if font is None:
+        font = get_default_font()
+    return PIL.ImageFont.truetype(font=font, size=size)
+
+#----------------------------------------------------------------------------
+
+def get_array(string, *, dropshadow_radius: int=None, **kwargs):
+    if dropshadow_radius is not None:
+        offset_x = int(np.ceil(dropshadow_radius*2/3))
+        offset_y = int(np.ceil(dropshadow_radius*2/3))
+        return _get_array_priv(string, dropshadow_radius=dropshadow_radius, offset_x=offset_x, offset_y=offset_y, **kwargs)
+    else:
+        return _get_array_priv(string, **kwargs)
+
+@functools.lru_cache(maxsize=10000)
+def _get_array_priv(
+    string: str, *,
+    size: int = 32,
+    max_width: Optional[int]=None,
+    max_height: Optional[int]=None,
+    min_size=10,
+    shrink_coef=0.8,
+    dropshadow_radius: int=None,
+    offset_x: int=None,
+    offset_y: int=None,
+    **kwargs
+):
+    cur_size = size
+    array = None
+    while True:
+        if dropshadow_radius is not None:
+            # separate implementation for dropshadow text rendering
+            array = _get_array_impl_dropshadow(string, size=cur_size, radius=dropshadow_radius, offset_x=offset_x, offset_y=offset_y, **kwargs)
+        else:
+            array = _get_array_impl(string, size=cur_size, **kwargs)
+        height, width, _ = array.shape
+        if (max_width is None or width <= max_width) and (max_height is None or height <= max_height) or (cur_size <= min_size):
+            break
+        cur_size = max(int(cur_size * shrink_coef), min_size)
+    return array
+
+#----------------------------------------------------------------------------
+
+@functools.lru_cache(maxsize=10000)
+def _get_array_impl(string, *, font=None, size=32, outline=0, outline_pad=3, outline_coef=3, outline_exp=2, line_pad: int=None):
+    pil_font = get_pil_font(font=font, size=size)
+    lines = [pil_font.getmask(line, 'L') for line in string.split('\n')]
+    lines = [np.array(line, dtype=np.uint8).reshape([line.size[1], line.size[0]]) for line in lines]
+    width = max(line.shape[1] for line in lines)
+    lines = [np.pad(line, ((0, 0), (0, width - line.shape[1])), mode='constant') for line in lines]
+    line_spacing = line_pad if line_pad is not None else size // 2
+    lines = [np.pad(line, ((0, line_spacing), (0, 0)), mode='constant') for line in lines[:-1]] + lines[-1:]
+    mask = np.concatenate(lines, axis=0)
+    alpha = mask
+    if outline > 0:
+        mask = np.pad(mask, int(np.ceil(outline * outline_pad)), mode='constant', constant_values=0)
+        alpha = mask.astype(np.float32) / 255
+        alpha = scipy.ndimage.gaussian_filter(alpha, outline)
+        alpha = 1 - np.maximum(1 - alpha * outline_coef, 0) ** outline_exp
+        alpha = (alpha * 255 + 0.5).clip(0, 255).astype(np.uint8)
+        alpha = np.maximum(alpha, mask)
+    return np.stack([mask, alpha], axis=-1)
+
+#----------------------------------------------------------------------------
+
+@functools.lru_cache(maxsize=10000)
+def _get_array_impl_dropshadow(string, *, font=None, size=32, radius: int, offset_x: int, offset_y: int, line_pad: int=None, **kwargs):
+    assert (offset_x > 0) and (offset_y > 0)
+    pil_font = get_pil_font(font=font, size=size)
+    lines = [pil_font.getmask(line, 'L') for line in string.split('\n')]
+    lines = [np.array(line, dtype=np.uint8).reshape([line.size[1], line.size[0]]) for line in lines]
+    width = max(line.shape[1] for line in lines)
+    lines = [np.pad(line, ((0, 0), (0, width - line.shape[1])), mode='constant') for line in lines]
+    line_spacing = line_pad if line_pad is not None else size // 2
+    lines = [np.pad(line, ((0, line_spacing), (0, 0)), mode='constant') for line in lines[:-1]] + lines[-1:]
+    mask = np.concatenate(lines, axis=0)
+    alpha = mask
+
+    mask = np.pad(mask, 2*radius + max(abs(offset_x), abs(offset_y)), mode='constant', constant_values=0)
+    alpha = mask.astype(np.float32) / 255
+    alpha = scipy.ndimage.gaussian_filter(alpha, radius)
+    alpha = 1 - np.maximum(1 - alpha * 1.5, 0) ** 1.4
+    alpha = (alpha * 255 + 0.5).clip(0, 255).astype(np.uint8)
+    alpha = np.pad(alpha, [(offset_y, 0), (offset_x, 0)], mode='constant')[:-offset_y, :-offset_x]
+    alpha = np.maximum(alpha, mask)
+    return np.stack([mask, alpha], axis=-1)
+
+#----------------------------------------------------------------------------
+
+@functools.lru_cache(maxsize=10000)
+def get_texture(string, bilinear=True, mipmap=True, **kwargs):
+    return gl_utils.Texture(image=get_array(string, **kwargs), bilinear=bilinear, mipmap=mipmap)
+
+#----------------------------------------------------------------------------

launcher.py

@@ -0,0 +1,189 @@
+#!/usr/bin/env python3 -u
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+import random, shlex, datetime
+import os, sys, subprocess, shutil
+from glob import iglob
+
+
+def copy_all_python_files(
+    source, snapshot_main_dir, code_snapshot_hash, recurse_dirs="fairseq"
+):
+    """
+    Copies following files from source to destination:
+        a) all *.py files at direct source location.
+        b) all fairseq/*.py recursively (default); recurse through comma-separated recurse_dirs
+    """
+    os.makedirs(snapshot_main_dir, exist_ok=True)
+    destination = os.path.join(snapshot_main_dir, code_snapshot_hash)
+    assert not os.path.exists(destination), "Code snapshot: {0} alredy exists".format(
+        code_snapshot_hash
+    )
+    os.makedirs(destination)
+
+    def all_pys(recurse_dirs):
+        yield from iglob(os.path.join(source, "*.py"))
+        for d in recurse_dirs.split(","):
+            yield from iglob(os.path.join(source, d, "**/*.py"), recursive=True)
+            yield from iglob(os.path.join(source, d, "**/*.so"), recursive=True)
+            yield from iglob(os.path.join(source, d, "**/*.yaml"), recursive=True)
+
+    for filepath in all_pys(recurse_dirs):
+        directory, filename = os.path.split(filepath)
+        if directory:
+            os.makedirs(os.path.join(destination, directory), exist_ok=True)
+        shutil.copy2(
+            os.path.join(source, filepath), os.path.join(destination, filepath)
+        )
+    return destination
+
+def launch_cluster(slurm_args, model_args):
+    # prepare
+    jobname = slurm_args.get('job-name', 'test')
+    if slurm_args.get('workplace') is not None:
+        os.makedirs(slurm_args.get('workplace'), exist_ok=True)
+    if slurm_args.get('workplace') is not None:
+        train_log = os.path.join(slurm_args['workplace'], 'train.%A.out')
+        train_stderr = os.path.join(slurm_args['workplace'], 'train.%A.stderr.%j')
+    else:
+        train_log = train_stderr = None
+    nodes, gpus = slurm_args.get('nodes', 1), slurm_args.get('gpus', 8)
+    if not slurm_args.get('local', False):
+        assert (train_log is not None) and (train_stderr is not None)
+    # parse slurm
+
+    destination = ""
+    # if slurm_args.get('workplace', None) is not None:
+    #     # Currently hash is just the current time in ISO format.
+    #     # Remove colons since they cannot be escaped in POSIX PATH env vars.
+    #     code_snapshot_hash = datetime.datetime.now().isoformat().replace(":", "_")
+    #     destination = copy_all_python_files(
+    #         ".",
+    #         os.path.join(slurm_args['workplace'], "slurm_snapshot_code"),
+    #         code_snapshot_hash,
+    #         'fairseq',
+    #     )
+    #     os.environ["PYTHONPATH"] = destination + ":" + os.environ.get("PYTHONPATH", "")
+    #     print('creat snapshot at {}'.format(destination))
+
+    train_cmd = ['python', os.path.join(destination, 'run_train.py'), ]
+    train_cmd.extend([f'gpus={nodes * gpus}'])
+    train_cmd.extend([f'port={get_random_port()}'])
+    train_cmd += model_args
+
+    base_srun_cmd = [
+            'srun',
+            '--job-name', jobname,
+            '--output', train_log,
+            '--error', train_stderr,
+            '--open-mode', 'append',
+            '--unbuffered',
+        ]
+    srun_cmd = base_srun_cmd + train_cmd
+    srun_cmd_str = ' '.join(map(shlex.quote, srun_cmd)) 
+    srun_cmd_str = srun_cmd_str + ' &'
+
+    sbatch_cmd = [
+                'sbatch',
+                '--job-name', jobname,
+                '--partition', slurm_args.get('partition', 'learnfair'),
+                '--gres', 'gpu:volta:{}'.format(gpus),
+                '--nodes', str(nodes),
+                '--ntasks-per-node', '1',
+                '--cpus-per-task', '20',
+                '--output', train_log,
+                '--error', train_stderr,
+                '--open-mode', 'append',
+                '--signal', 'B:USR1@180',
+                '--time', slurm_args.get('time', '4320'),
+                '--mem', slurm_args.get('mem', '500gb'),
+                '--exclusive',
+                '--exclude', 'learnfair5035,learnfair5289,learnfair5088,learnfair5028,learnfair5032,learnfair5033,learnfair5056,learnfair5098,learnfair5122,learnfair5124,learnfair5156,learnfair5036,learnfair5258,learnfair5205,learnfair5201,learnfair5240,learnfair5087,learnfair5119,learnfair5246,learnfair7474,learnfair7585,learnfair5150,learnfair5166,learnfair5215,learnfair5142,learnfair5070,learnfair5236,learnfair7523'
+            ]
+    if 'constraint' in slurm_args:
+        sbatch_cmd += ['-C', slurm_args.get('constraint')]
+    if 'comment' in slurm_args:
+        sbatch_cmd += ['--comment', slurm_args.get('comment')]
+    
+    wrapped_cmd = requeue_support() + '\n' + srun_cmd_str + ' \n wait $! \n sleep 610 & \n wait $!'
+    sbatch_cmd += ['--wrap', wrapped_cmd]
+    sbatch_cmd_str = ' '.join(map(shlex.quote, sbatch_cmd))
+    
+    # start training
+    env = os.environ.copy()
+    env['OMP_NUM_THREADS'] = '2'
+    env['NCCL_SOCKET_IFNAME'] = ''
+    
+    if env.get('SLURM_ARGS', None) is not None:
+        del env['SLURM_ARGS']
+
+    if nodes > 1:
+        env['NCCL_SOCKET_IFNAME'] = '^docker0,lo'
+        env['NCCL_DEBUG'] = 'INFO'
+
+    if slurm_args.get('dry-run', False):
+        print(sbatch_cmd_str)
+    
+    elif slurm_args.get('local', False):
+        assert nodes == 1, 'distributed training cannot be combined with local' 
+        if 'CUDA_VISIBLE_DEVICES' not in env:
+            env['CUDA_VISIBLE_DEVICES'] = ','.join(map(str, range(gpus)))
+        env['NCCL_DEBUG'] = 'INFO'
+        
+        if train_log is not None:
+            train_proc = subprocess.Popen(train_cmd, env=env, stdout=subprocess.PIPE)
+            tee_proc = subprocess.Popen(['tee', '-a', train_log], stdin=train_proc.stdout)
+            train_proc.stdout.close()
+            train_proc.wait()
+            tee_proc.wait()
+        else:
+            train_proc = subprocess.Popen(train_cmd, env=env)
+            train_proc.wait()
+    else:
+        with open(train_log, 'a') as train_log_h:
+            print(f'running command: {sbatch_cmd_str}\n')
+            with subprocess.Popen(sbatch_cmd, stdout=subprocess.PIPE, env=env) as train_proc:
+                stdout = train_proc.stdout.read().decode('utf-8')
+                print(stdout, file=train_log_h)
+                try:
+                    job_id = int(stdout.rstrip().split()[-1])
+                    return job_id
+                except IndexError:
+                    return None
+
+
+def launch(slurm_args, model_args):
+    job_id = launch_cluster(slurm_args, model_args)
+    if job_id is not None:
+        print('Launched {}'.format(job_id))
+    else:
+        print('Failed.')
+
+
+def requeue_support():
+    return """
+        trap_handler () {
+           echo "Caught signal: " $1
+           # SIGTERM must be bypassed
+           if [ "$1" = "TERM" ]; then
+               echo "bypass sigterm"
+           else
+             # Submit a new job to the queue
+             echo "Requeuing " $SLURM_JOB_ID
+             scontrol requeue $SLURM_JOB_ID
+           fi
+        }
+
+
+        # Install signal handler
+        trap 'trap_handler USR1' USR1
+        trap 'trap_handler TERM' TERM
+    """
+
+
+def get_random_port():
+    old_state = random.getstate()
+    random.seed()
+    port = random.randint(10000, 20000)
+    random.setstate(old_state)
+    return port

legacy.py

@@ -0,0 +1,320 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import click
+import pickle
+import re
+import copy
+import numpy as np
+import torch
+import dnnlib
+from torch_utils import misc
+
+#----------------------------------------------------------------------------
+
+def load_network_pkl(f, force_fp16=False):
+    data = _LegacyUnpickler(f).load()
+
+    # Legacy TensorFlow pickle => convert.
+    if isinstance(data, tuple) and len(data) == 3 and all(isinstance(net, _TFNetworkStub) for net in data):
+        tf_G, tf_D, tf_Gs = data
+        G = convert_tf_generator(tf_G)
+        D = convert_tf_discriminator(tf_D)
+        G_ema = convert_tf_generator(tf_Gs)
+        data = dict(G=G, D=D, G_ema=G_ema)
+
+    # Add missing fields.
+    if 'training_set_kwargs' not in data:
+        data['training_set_kwargs'] = None
+    if 'augment_pipe' not in data:
+        data['augment_pipe'] = None
+
+    # Validate contents.
+    # assert isinstance(data['G'], torch.nn.Module)
+    # assert isinstance(data['D'], torch.nn.Module)
+    # assert isinstance(data['G_ema'], torch.nn.Module)
+    # assert isinstance(data['training_set_kwargs'], (dict, type(None)))
+    # assert isinstance(data['augment_pipe'], (torch.nn.Module, type(None)))
+
+    # Force FP16.
+    if force_fp16:
+        for key in ['G', 'D', 'G_ema']:
+            old = data[key]
+            kwargs = copy.deepcopy(old.init_kwargs)
+            if key.startswith('G'):
+                kwargs.synthesis_kwargs = dnnlib.EasyDict(kwargs.get('synthesis_kwargs', {}))
+                kwargs.synthesis_kwargs.num_fp16_res = 4
+                kwargs.synthesis_kwargs.conv_clamp = 256
+            if key.startswith('D'):
+                kwargs.num_fp16_res = 4
+                kwargs.conv_clamp = 256
+            if kwargs != old.init_kwargs:
+                new = type(old)(**kwargs).eval().requires_grad_(False)
+                misc.copy_params_and_buffers(old, new, require_all=True)
+                data[key] = new
+    return data
+
+#----------------------------------------------------------------------------
+
+class _TFNetworkStub(dnnlib.EasyDict):
+    pass
+
+class _LegacyUnpickler(pickle.Unpickler):
+    def find_class(self, module, name):
+        if module == 'dnnlib.tflib.network' and name == 'Network':
+            return _TFNetworkStub
+        return super().find_class(module, name)
+
+#----------------------------------------------------------------------------
+
+def _collect_tf_params(tf_net):
+    # pylint: disable=protected-access
+    tf_params = dict()
+    def recurse(prefix, tf_net):
+        for name, value in tf_net.variables:
+            tf_params[prefix + name] = value
+        for name, comp in tf_net.components.items():
+            recurse(prefix + name + '/', comp)
+    recurse('', tf_net)
+    return tf_params
+
+#----------------------------------------------------------------------------
+
+def _populate_module_params(module, *patterns):
+    for name, tensor in misc.named_params_and_buffers(module):
+        found = False
+        value = None
+        for pattern, value_fn in zip(patterns[0::2], patterns[1::2]):
+            match = re.fullmatch(pattern, name)
+            if match:
+                found = True
+                if value_fn is not None:
+                    value = value_fn(*match.groups())
+                break
+        try:
+            assert found
+            if value is not None:
+                tensor.copy_(torch.from_numpy(np.array(value)))
+        except:
+            print(name, list(tensor.shape))
+            raise
+
+#----------------------------------------------------------------------------
+
+def convert_tf_generator(tf_G):
+    if tf_G.version < 4:
+        raise ValueError('TensorFlow pickle version too low')
+
+    # Collect kwargs.
+    tf_kwargs = tf_G.static_kwargs
+    known_kwargs = set()
+    def kwarg(tf_name, default=None, none=None):
+        known_kwargs.add(tf_name)
+        val = tf_kwargs.get(tf_name, default)
+        return val if val is not None else none
+
+    # Convert kwargs.
+    kwargs = dnnlib.EasyDict(
+        z_dim                   = kwarg('latent_size',          512),
+        c_dim                   = kwarg('label_size',           0),
+        w_dim                   = kwarg('dlatent_size',         512),
+        img_resolution          = kwarg('resolution',           1024),
+        img_channels            = kwarg('num_channels',         3),
+        mapping_kwargs = dnnlib.EasyDict(
+            num_layers          = kwarg('mapping_layers',       8),
+            embed_features      = kwarg('label_fmaps',          None),
+            layer_features      = kwarg('mapping_fmaps',        None),
+            activation          = kwarg('mapping_nonlinearity', 'lrelu'),
+            lr_multiplier       = kwarg('mapping_lrmul',        0.01),
+            w_avg_beta          = kwarg('w_avg_beta',           0.995,  none=1),
+        ),
+        synthesis_kwargs = dnnlib.EasyDict(
+            channel_base        = kwarg('fmap_base',            16384) * 2,
+            channel_max         = kwarg('fmap_max',             512),
+            num_fp16_res        = kwarg('num_fp16_res',         0),
+            conv_clamp          = kwarg('conv_clamp',           None),
+            architecture        = kwarg('architecture',         'skip'),
+            resample_filter     = kwarg('resample_kernel',      [1,3,3,1]),
+            use_noise           = kwarg('use_noise',            True),
+            activation          = kwarg('nonlinearity',         'lrelu'),
+        ),
+    )
+
+    # Check for unknown kwargs.
+    kwarg('truncation_psi')
+    kwarg('truncation_cutoff')
+    kwarg('style_mixing_prob')
+    kwarg('structure')
+    unknown_kwargs = list(set(tf_kwargs.keys()) - known_kwargs)
+    if len(unknown_kwargs) > 0:
+        raise ValueError('Unknown TensorFlow kwarg', unknown_kwargs[0])
+
+    # Collect params.
+    tf_params = _collect_tf_params(tf_G)
+    for name, value in list(tf_params.items()):
+        match = re.fullmatch(r'ToRGB_lod(\d+)/(.*)', name)
+        if match:
+            r = kwargs.img_resolution // (2 ** int(match.group(1)))
+            tf_params[f'{r}x{r}/ToRGB/{match.group(2)}'] = value
+            kwargs.synthesis.kwargs.architecture = 'orig'
+    #for name, value in tf_params.items(): print(f'{name:<50s}{list(value.shape)}')
+
+    # Convert params.
+    from training import networks
+    G = networks.Generator(**kwargs).eval().requires_grad_(False)
+    # pylint: disable=unnecessary-lambda
+    _populate_module_params(G,
+        r'mapping\.w_avg',                                  lambda:     tf_params[f'dlatent_avg'],
+        r'mapping\.embed\.weight',                          lambda:     tf_params[f'mapping/LabelEmbed/weight'].transpose(),
+        r'mapping\.embed\.bias',                            lambda:     tf_params[f'mapping/LabelEmbed/bias'],
+        r'mapping\.fc(\d+)\.weight',                        lambda i:   tf_params[f'mapping/Dense{i}/weight'].transpose(),
+        r'mapping\.fc(\d+)\.bias',                          lambda i:   tf_params[f'mapping/Dense{i}/bias'],
+        r'synthesis\.b4\.const',                            lambda:     tf_params[f'synthesis/4x4/Const/const'][0],
+        r'synthesis\.b4\.conv1\.weight',                    lambda:     tf_params[f'synthesis/4x4/Conv/weight'].transpose(3, 2, 0, 1),
+        r'synthesis\.b4\.conv1\.bias',                      lambda:     tf_params[f'synthesis/4x4/Conv/bias'],
+        r'synthesis\.b4\.conv1\.noise_const',               lambda:     tf_params[f'synthesis/noise0'][0, 0],
+        r'synthesis\.b4\.conv1\.noise_strength',            lambda:     tf_params[f'synthesis/4x4/Conv/noise_strength'],
+        r'synthesis\.b4\.conv1\.affine\.weight',            lambda:     tf_params[f'synthesis/4x4/Conv/mod_weight'].transpose(),
+        r'synthesis\.b4\.conv1\.affine\.bias',              lambda:     tf_params[f'synthesis/4x4/Conv/mod_bias'] + 1,
+        r'synthesis\.b(\d+)\.conv0\.weight',                lambda r:   tf_params[f'synthesis/{r}x{r}/Conv0_up/weight'][::-1, ::-1].transpose(3, 2, 0, 1),
+        r'synthesis\.b(\d+)\.conv0\.bias',                  lambda r:   tf_params[f'synthesis/{r}x{r}/Conv0_up/bias'],
+        r'synthesis\.b(\d+)\.conv0\.noise_const',           lambda r:   tf_params[f'synthesis/noise{int(np.log2(int(r)))*2-5}'][0, 0],
+        r'synthesis\.b(\d+)\.conv0\.noise_strength',        lambda r:   tf_params[f'synthesis/{r}x{r}/Conv0_up/noise_strength'],
+        r'synthesis\.b(\d+)\.conv0\.affine\.weight',        lambda r:   tf_params[f'synthesis/{r}x{r}/Conv0_up/mod_weight'].transpose(),
+        r'synthesis\.b(\d+)\.conv0\.affine\.bias',          lambda r:   tf_params[f'synthesis/{r}x{r}/Conv0_up/mod_bias'] + 1,
+        r'synthesis\.b(\d+)\.conv1\.weight',                lambda r:   tf_params[f'synthesis/{r}x{r}/Conv1/weight'].transpose(3, 2, 0, 1),
+        r'synthesis\.b(\d+)\.conv1\.bias',                  lambda r:   tf_params[f'synthesis/{r}x{r}/Conv1/bias'],
+        r'synthesis\.b(\d+)\.conv1\.noise_const',           lambda r:   tf_params[f'synthesis/noise{int(np.log2(int(r)))*2-4}'][0, 0],
+        r'synthesis\.b(\d+)\.conv1\.noise_strength',        lambda r:   tf_params[f'synthesis/{r}x{r}/Conv1/noise_strength'],
+        r'synthesis\.b(\d+)\.conv1\.affine\.weight',        lambda r:   tf_params[f'synthesis/{r}x{r}/Conv1/mod_weight'].transpose(),
+        r'synthesis\.b(\d+)\.conv1\.affine\.bias',          lambda r:   tf_params[f'synthesis/{r}x{r}/Conv1/mod_bias'] + 1,
+        r'synthesis\.b(\d+)\.torgb\.weight',                lambda r:   tf_params[f'synthesis/{r}x{r}/ToRGB/weight'].transpose(3, 2, 0, 1),
+        r'synthesis\.b(\d+)\.torgb\.bias',                  lambda r:   tf_params[f'synthesis/{r}x{r}/ToRGB/bias'],
+        r'synthesis\.b(\d+)\.torgb\.affine\.weight',        lambda r:   tf_params[f'synthesis/{r}x{r}/ToRGB/mod_weight'].transpose(),
+        r'synthesis\.b(\d+)\.torgb\.affine\.bias',          lambda r:   tf_params[f'synthesis/{r}x{r}/ToRGB/mod_bias'] + 1,
+        r'synthesis\.b(\d+)\.skip\.weight',                 lambda r:   tf_params[f'synthesis/{r}x{r}/Skip/weight'][::-1, ::-1].transpose(3, 2, 0, 1),
+        r'.*\.resample_filter',                             None,
+    )
+    return G
+
+#----------------------------------------------------------------------------
+
+def convert_tf_discriminator(tf_D):
+    if tf_D.version < 4:
+        raise ValueError('TensorFlow pickle version too low')
+
+    # Collect kwargs.
+    tf_kwargs = tf_D.static_kwargs
+    known_kwargs = set()
+    def kwarg(tf_name, default=None):
+        known_kwargs.add(tf_name)
+        return tf_kwargs.get(tf_name, default)
+
+    # Convert kwargs.
+    kwargs = dnnlib.EasyDict(
+        c_dim                   = kwarg('label_size',           0),
+        img_resolution          = kwarg('resolution',           1024),
+        img_channels            = kwarg('num_channels',         3),
+        architecture            = kwarg('architecture',         'resnet'),
+        channel_base            = kwarg('fmap_base',            16384) * 2,
+        channel_max             = kwarg('fmap_max',             512),
+        num_fp16_res            = kwarg('num_fp16_res',         0),
+        conv_clamp              = kwarg('conv_clamp',           None),
+        cmap_dim                = kwarg('mapping_fmaps',        None),
+        block_kwargs = dnnlib.EasyDict(
+            activation          = kwarg('nonlinearity',         'lrelu'),
+            resample_filter     = kwarg('resample_kernel',      [1,3,3,1]),
+            freeze_layers       = kwarg('freeze_layers',        0),
+        ),
+        mapping_kwargs = dnnlib.EasyDict(
+            num_layers          = kwarg('mapping_layers',       0),
+            embed_features      = kwarg('mapping_fmaps',        None),
+            layer_features      = kwarg('mapping_fmaps',        None),
+            activation          = kwarg('nonlinearity',         'lrelu'),
+            lr_multiplier       = kwarg('mapping_lrmul',        0.1),
+        ),
+        epilogue_kwargs = dnnlib.EasyDict(
+            mbstd_group_size    = kwarg('mbstd_group_size',     None),
+            mbstd_num_channels  = kwarg('mbstd_num_features',   1),
+            activation          = kwarg('nonlinearity',         'lrelu'),
+        ),
+    )
+
+    # Check for unknown kwargs.
+    kwarg('structure')
+    unknown_kwargs = list(set(tf_kwargs.keys()) - known_kwargs)
+    if len(unknown_kwargs) > 0:
+        raise ValueError('Unknown TensorFlow kwarg', unknown_kwargs[0])
+
+    # Collect params.
+    tf_params = _collect_tf_params(tf_D)
+    for name, value in list(tf_params.items()):
+        match = re.fullmatch(r'FromRGB_lod(\d+)/(.*)', name)
+        if match:
+            r = kwargs.img_resolution // (2 ** int(match.group(1)))
+            tf_params[f'{r}x{r}/FromRGB/{match.group(2)}'] = value
+            kwargs.architecture = 'orig'
+    #for name, value in tf_params.items(): print(f'{name:<50s}{list(value.shape)}')
+
+    # Convert params.
+    from training import networks
+    D = networks.Discriminator(**kwargs).eval().requires_grad_(False)
+    # pylint: disable=unnecessary-lambda
+    _populate_module_params(D,
+        r'b(\d+)\.fromrgb\.weight',     lambda r:       tf_params[f'{r}x{r}/FromRGB/weight'].transpose(3, 2, 0, 1),
+        r'b(\d+)\.fromrgb\.bias',       lambda r:       tf_params[f'{r}x{r}/FromRGB/bias'],
+        r'b(\d+)\.conv(\d+)\.weight',   lambda r, i:    tf_params[f'{r}x{r}/Conv{i}{["","_down"][int(i)]}/weight'].transpose(3, 2, 0, 1),
+        r'b(\d+)\.conv(\d+)\.bias',     lambda r, i:    tf_params[f'{r}x{r}/Conv{i}{["","_down"][int(i)]}/bias'],
+        r'b(\d+)\.skip\.weight',        lambda r:       tf_params[f'{r}x{r}/Skip/weight'].transpose(3, 2, 0, 1),
+        r'mapping\.embed\.weight',      lambda:         tf_params[f'LabelEmbed/weight'].transpose(),
+        r'mapping\.embed\.bias',        lambda:         tf_params[f'LabelEmbed/bias'],
+        r'mapping\.fc(\d+)\.weight',    lambda i:       tf_params[f'Mapping{i}/weight'].transpose(),
+        r'mapping\.fc(\d+)\.bias',      lambda i:       tf_params[f'Mapping{i}/bias'],
+        r'b4\.conv\.weight',            lambda:         tf_params[f'4x4/Conv/weight'].transpose(3, 2, 0, 1),
+        r'b4\.conv\.bias',              lambda:         tf_params[f'4x4/Conv/bias'],
+        r'b4\.fc\.weight',              lambda:         tf_params[f'4x4/Dense0/weight'].transpose(),
+        r'b4\.fc\.bias',                lambda:         tf_params[f'4x4/Dense0/bias'],
+        r'b4\.out\.weight',             lambda:         tf_params[f'Output/weight'].transpose(),
+        r'b4\.out\.bias',               lambda:         tf_params[f'Output/bias'],
+        r'.*\.resample_filter',         None,
+    )
+    return D
+
+#----------------------------------------------------------------------------
+
+@click.command()
+@click.option('--source', help='Input pickle', required=True, metavar='PATH')
+@click.option('--dest', help='Output pickle', required=True, metavar='PATH')
+@click.option('--force-fp16', help='Force the networks to use FP16', type=bool, default=False, metavar='BOOL', show_default=True)
+def convert_network_pickle(source, dest, force_fp16):
+    """Convert legacy network pickle into the native PyTorch format.
+
+    The tool is able to load the main network configurations exported using the TensorFlow version of StyleGAN2 or StyleGAN2-ADA.
+    It does not support e.g. StyleGAN2-ADA comparison methods, StyleGAN2 configs A-D, or StyleGAN1 networks.
+
+    Example:
+
+    \b
+    python legacy.py \\
+        --source=https://nvlabs-fi-cdn.nvidia.com/stylegan2/networks/stylegan2-cat-config-f.pkl \\
+        --dest=stylegan2-cat-config-f.pkl
+    """
+    print(f'Loading "{source}"...')
+    with dnnlib.util.open_url(source) as f:
+        data = load_network_pkl(f, force_fp16=force_fp16)
+    print(f'Saving "{dest}"...')
+    with open(dest, 'wb') as f:
+        pickle.dump(data, f)
+    print('Done.')
+
+#----------------------------------------------------------------------------
+
+if __name__ == "__main__":
+    convert_network_pickle() # pylint: disable=no-value-for-parameter
+
+#----------------------------------------------------------------------------

metrics/__init__.py

@@ -0,0 +1,9 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+# empty

metrics/frechet_inception_distance.py

@@ -0,0 +1,40 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+"""Frechet Inception Distance (FID) from the paper
+"GANs trained by a two time-scale update rule converge to a local Nash
+equilibrium". Matches the original implementation by Heusel et al. at
+https://github.com/bioinf-jku/TTUR/blob/master/fid.py"""
+
+import numpy as np
+import scipy.linalg
+from . import metric_utils
+
+#----------------------------------------------------------------------------
+
+def compute_fid(opts, max_real, num_gen):
+    # Direct TorchScript translation of http://download.tensorflow.org/models/image/imagenet/inception-2015-12-05.tgz
+    detector_url = 'https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/metrics/inception-2015-12-05.pt'
+    detector_kwargs = dict(return_features=True) # Return raw features before the softmax layer.
+    mu_real, sigma_real = metric_utils.compute_feature_stats_for_dataset(
+        opts=opts, detector_url=detector_url, detector_kwargs=detector_kwargs,
+        rel_lo=0, rel_hi=0, capture_mean_cov=True, max_items=max_real).get_mean_cov()
+
+    mu_gen, sigma_gen = metric_utils.compute_feature_stats_for_generator(
+        opts=opts, detector_url=detector_url, detector_kwargs=detector_kwargs,
+        rel_lo=0, rel_hi=1, capture_mean_cov=True, max_items=num_gen).get_mean_cov()
+
+    if opts.rank != 0:
+        return float('nan')
+
+    m = np.square(mu_gen - mu_real).sum()
+    s, _ = scipy.linalg.sqrtm(np.dot(sigma_gen, sigma_real), disp=False) # pylint: disable=no-member
+    fid = np.real(m + np.trace(sigma_gen + sigma_real - s * 2))
+    return float(fid)
+
+#----------------------------------------------------------------------------

metrics/inception_score.py

@@ -0,0 +1,38 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+"""Inception Score (IS) from the paper "Improved techniques for training
+GANs". Matches the original implementation by Salimans et al. at
+https://github.com/openai/improved-gan/blob/master/inception_score/model.py"""
+
+import numpy as np
+from . import metric_utils
+
+#----------------------------------------------------------------------------
+
+def compute_is(opts, num_gen, num_splits):
+    # Direct TorchScript translation of http://download.tensorflow.org/models/image/imagenet/inception-2015-12-05.tgz
+    detector_url = 'https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/metrics/inception-2015-12-05.pt'
+    detector_kwargs = dict(no_output_bias=True) # Match the original implementation by not applying bias in the softmax layer.
+
+    gen_probs = metric_utils.compute_feature_stats_for_generator(
+        opts=opts, detector_url=detector_url, detector_kwargs=detector_kwargs,
+        capture_all=True, max_items=num_gen).get_all()
+
+    if opts.rank != 0:
+        return float('nan'), float('nan')
+
+    scores = []
+    for i in range(num_splits):
+        part = gen_probs[i * num_gen // num_splits : (i + 1) * num_gen // num_splits]
+        kl = part * (np.log(part) - np.log(np.mean(part, axis=0, keepdims=True)))
+        kl = np.mean(np.sum(kl, axis=1))
+        scores.append(np.exp(kl))
+    return float(np.mean(scores)), float(np.std(scores))
+
+#----------------------------------------------------------------------------

metrics/kernel_inception_distance.py

@@ -0,0 +1,46 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+"""Kernel Inception Distance (KID) from the paper "Demystifying MMD
+GANs". Matches the original implementation by Binkowski et al. at
+https://github.com/mbinkowski/MMD-GAN/blob/master/gan/compute_scores.py"""
+
+import numpy as np
+from . import metric_utils
+
+#----------------------------------------------------------------------------
+
+def compute_kid(opts, max_real, num_gen, num_subsets, max_subset_size):
+    # Direct TorchScript translation of http://download.tensorflow.org/models/image/imagenet/inception-2015-12-05.tgz
+    detector_url = 'https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/metrics/inception-2015-12-05.pt'
+    detector_kwargs = dict(return_features=True) # Return raw features before the softmax layer.
+
+    real_features = metric_utils.compute_feature_stats_for_dataset(
+        opts=opts, detector_url=detector_url, detector_kwargs=detector_kwargs,
+        rel_lo=0, rel_hi=0, capture_all=True, max_items=max_real).get_all()
+
+    gen_features = metric_utils.compute_feature_stats_for_generator(
+        opts=opts, detector_url=detector_url, detector_kwargs=detector_kwargs,
+        rel_lo=0, rel_hi=1, capture_all=True, max_items=num_gen).get_all()
+
+    if opts.rank != 0:
+        return float('nan')
+
+    n = real_features.shape[1]
+    m = min(min(real_features.shape[0], gen_features.shape[0]), max_subset_size)
+    t = 0
+    for _subset_idx in range(num_subsets):
+        x = gen_features[np.random.choice(gen_features.shape[0], m, replace=False)]
+        y = real_features[np.random.choice(real_features.shape[0], m, replace=False)]
+        a = (x @ x.T / n + 1) ** 3 + (y @ y.T / n + 1) ** 3
+        b = (x @ y.T / n + 1) ** 3
+        t += (a.sum() - np.diag(a).sum()) / (m - 1) - b.sum() * 2 / m
+    kid = t / num_subsets / m
+    return float(kid)
+
+#----------------------------------------------------------------------------

metrics/metric_main.py

@@ -0,0 +1,152 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import os
+import time
+import json
+import torch
+import dnnlib
+
+from . import metric_utils
+from . import frechet_inception_distance
+from . import kernel_inception_distance
+from . import precision_recall
+from . import perceptual_path_length
+from . import inception_score
+
+#----------------------------------------------------------------------------
+
+_metric_dict = dict() # name => fn
+
+def register_metric(fn):
+    assert callable(fn)
+    _metric_dict[fn.__name__] = fn
+    return fn
+
+def is_valid_metric(metric):
+    return metric in _metric_dict
+
+def list_valid_metrics():
+    return list(_metric_dict.keys())
+
+#----------------------------------------------------------------------------
+
+def calc_metric(metric, **kwargs): # See metric_utils.MetricOptions for the full list of arguments.
+    assert is_valid_metric(metric)
+    opts = metric_utils.MetricOptions(**kwargs)
+
+    # Calculate.
+    start_time = time.time()
+    results = _metric_dict[metric](opts)
+    total_time = time.time() - start_time
+
+    # Broadcast results.
+    for key, value in list(results.items()):
+        if opts.num_gpus > 1:
+            value = torch.as_tensor(value, dtype=torch.float64, device=opts.device)
+            torch.distributed.broadcast(tensor=value, src=0)
+            value = float(value.cpu())
+        results[key] = value
+
+    # Decorate with metadata.
+    return dnnlib.EasyDict(
+        results         = dnnlib.EasyDict(results),
+        metric          = metric,
+        total_time      = total_time,
+        total_time_str  = dnnlib.util.format_time(total_time),
+        num_gpus        = opts.num_gpus,
+    )
+
+#----------------------------------------------------------------------------
+
+def report_metric(result_dict, run_dir=None, snapshot_pkl=None):
+    metric = result_dict['metric']
+    assert is_valid_metric(metric)
+    if run_dir is not None and snapshot_pkl is not None:
+        snapshot_pkl = os.path.relpath(snapshot_pkl, run_dir)
+
+    jsonl_line = json.dumps(dict(result_dict, snapshot_pkl=snapshot_pkl, timestamp=time.time()))
+    print(jsonl_line)
+    if run_dir is not None and os.path.isdir(run_dir):
+        with open(os.path.join(run_dir, f'metric-{metric}.jsonl'), 'at') as f:
+            f.write(jsonl_line + '\n')
+
+#----------------------------------------------------------------------------
+# Primary metrics.
+
+@register_metric
+def fid50k_full(opts):
+    opts.dataset_kwargs.update(max_size=None, xflip=False)
+    fid = frechet_inception_distance.compute_fid(opts, max_real=None, num_gen=50000)
+    return dict(fid50k_full=fid)
+
+@register_metric
+def kid50k_full(opts):
+    opts.dataset_kwargs.update(max_size=None, xflip=False)
+    kid = kernel_inception_distance.compute_kid(opts, max_real=1000000, num_gen=50000, num_subsets=100, max_subset_size=1000)
+    return dict(kid50k_full=kid)
+
+@register_metric
+def pr50k3_full(opts):
+    opts.dataset_kwargs.update(max_size=None, xflip=False)
+    precision, recall = precision_recall.compute_pr(opts, max_real=200000, num_gen=50000, nhood_size=3, row_batch_size=10000, col_batch_size=10000)
+    return dict(pr50k3_full_precision=precision, pr50k3_full_recall=recall)
+
+@register_metric
+def ppl2_wend(opts):
+    ppl = perceptual_path_length.compute_ppl(opts, num_samples=50000, epsilon=1e-4, space='w', sampling='end', crop=False, batch_size=2)
+    return dict(ppl2_wend=ppl)
+
+@register_metric
+def is50k(opts):
+    opts.dataset_kwargs.update(max_size=None, xflip=False)
+    mean, std = inception_score.compute_is(opts, num_gen=50000, num_splits=10)
+    return dict(is50k_mean=mean, is50k_std=std)
+
+#----------------------------------------------------------------------------
+# Legacy metrics.
+
+@register_metric
+def fid50k(opts):
+    opts.dataset_kwargs.update(max_size=None)
+    fid = frechet_inception_distance.compute_fid(opts, max_real=50000, num_gen=50000)
+    return dict(fid50k=fid)
+
+@register_metric
+def kid50k(opts):
+    opts.dataset_kwargs.update(max_size=None)
+    kid = kernel_inception_distance.compute_kid(opts, max_real=50000, num_gen=50000, num_subsets=100, max_subset_size=1000)
+    return dict(kid50k=kid)
+
+@register_metric
+def pr50k3(opts):
+    opts.dataset_kwargs.update(max_size=None)
+    precision, recall = precision_recall.compute_pr(opts, max_real=50000, num_gen=50000, nhood_size=3, row_batch_size=10000, col_batch_size=10000)
+    return dict(pr50k3_precision=precision, pr50k3_recall=recall)
+
+@register_metric
+def ppl_zfull(opts):
+    ppl = perceptual_path_length.compute_ppl(opts, num_samples=50000, epsilon=1e-4, space='z', sampling='full', crop=True, batch_size=2)
+    return dict(ppl_zfull=ppl)
+
+@register_metric
+def ppl_wfull(opts):
+    ppl = perceptual_path_length.compute_ppl(opts, num_samples=50000, epsilon=1e-4, space='w', sampling='full', crop=True, batch_size=2)
+    return dict(ppl_wfull=ppl)
+
+@register_metric
+def ppl_zend(opts):
+    ppl = perceptual_path_length.compute_ppl(opts, num_samples=50000, epsilon=1e-4, space='z', sampling='end', crop=True, batch_size=2)
+    return dict(ppl_zend=ppl)
+
+@register_metric
+def ppl_wend(opts):
+    ppl = perceptual_path_length.compute_ppl(opts, num_samples=50000, epsilon=1e-4, space='w', sampling='end', crop=True, batch_size=2)
+    return dict(ppl_wend=ppl)
+
+#----------------------------------------------------------------------------

metrics/metric_utils.py

@@ -0,0 +1,305 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import os
+import time
+import hashlib
+import pickle
+import copy
+import uuid
+import numpy as np
+import torch
+import dnnlib
+import glob
+#----------------------------------------------------------------------------
+
+class MetricOptions:
+    def __init__(self, G=None, G_kwargs={}, dataset_kwargs={}, num_gpus=1, rank=0, device=None, progress=None, cache=True):
+        assert 0 <= rank < num_gpus
+        self.G              = G
+        self.G_kwargs       = dnnlib.EasyDict(G_kwargs)
+        self.dataset_kwargs = dnnlib.EasyDict(dataset_kwargs)
+        self.num_gpus       = num_gpus
+        self.rank           = rank
+        self.device         = device if device is not None else torch.device('cuda', rank)
+        self.progress       = progress.sub() if progress is not None and rank == 0 else ProgressMonitor()
+        self.cache          = cache
+
+#----------------------------------------------------------------------------
+
+_feature_detector_cache = dict()
+
+def get_feature_detector_name(url):
+    return os.path.splitext(url.split('/')[-1])[0]
+
+def get_feature_detector(url, device=torch.device('cpu'), num_gpus=1, rank=0, verbose=False):
+    assert 0 <= rank < num_gpus
+    key = (url, device)
+    if key not in _feature_detector_cache:
+        is_leader = (rank == 0)
+        if not is_leader and num_gpus > 1:
+            torch.distributed.barrier() # leader goes first
+        with dnnlib.util.open_url(url, verbose=(verbose and is_leader)) as f:
+            _feature_detector_cache[key] = torch.jit.load(f).eval().to(device)
+        if is_leader and num_gpus > 1:
+            torch.distributed.barrier() # others follow
+    return _feature_detector_cache[key]
+
+#----------------------------------------------------------------------------
+
+class FeatureStats:
+    def __init__(self, capture_all=False, capture_mean_cov=False, max_items=None):
+        self.capture_all = capture_all
+        self.capture_mean_cov = capture_mean_cov
+        self.max_items = max_items
+        self.num_items = 0
+        self.num_features = None
+        self.all_features = None
+        self.raw_mean = None
+        self.raw_cov = None
+
+    def set_num_features(self, num_features):
+        if self.num_features is not None:
+            assert num_features == self.num_features
+        else:
+            self.num_features = num_features
+            self.all_features = []
+            self.raw_mean = np.zeros([num_features], dtype=np.float64)
+            self.raw_cov = np.zeros([num_features, num_features], dtype=np.float64)
+
+    def is_full(self):
+        return (self.max_items is not None) and (self.num_items >= self.max_items)
+
+    def append(self, x):
+        x = np.asarray(x, dtype=np.float32)
+        assert x.ndim == 2
+        if (self.max_items is not None) and (self.num_items + x.shape[0] > self.max_items):
+            if self.num_items >= self.max_items:
+                return
+            x = x[:self.max_items - self.num_items]
+
+        self.set_num_features(x.shape[1])
+        self.num_items += x.shape[0]
+        if self.capture_all:
+            self.all_features.append(x)
+        if self.capture_mean_cov:
+            x64 = x.astype(np.float64)
+            self.raw_mean += x64.sum(axis=0)
+            self.raw_cov += x64.T @ x64
+
+    def append_torch(self, x, num_gpus=1, rank=0):
+        assert isinstance(x, torch.Tensor) and x.ndim == 2
+        assert 0 <= rank < num_gpus
+        if num_gpus > 1:
+            ys = []
+            for src in range(num_gpus):
+                y = x.clone()
+                torch.distributed.broadcast(y, src=src)
+                ys.append(y)
+            x = torch.stack(ys, dim=1).flatten(0, 1) # interleave samples
+        self.append(x.cpu().numpy())
+
+    def get_all(self):
+        assert self.capture_all
+        return np.concatenate(self.all_features, axis=0)
+
+    def get_all_torch(self):
+        return torch.from_numpy(self.get_all())
+
+    def get_mean_cov(self):
+        assert self.capture_mean_cov
+        mean = self.raw_mean / self.num_items
+        cov = self.raw_cov / self.num_items
+        cov = cov - np.outer(mean, mean)
+        return mean, cov
+
+    def save(self, pkl_file):
+        with open(pkl_file, 'wb') as f:
+            pickle.dump(self.__dict__, f)
+
+    @staticmethod
+    def load(pkl_file):
+        with open(pkl_file, 'rb') as f:
+            s = dnnlib.EasyDict(pickle.load(f))
+        obj = FeatureStats(capture_all=s.capture_all, max_items=s.max_items)
+        obj.__dict__.update(s)
+        return obj
+
+#----------------------------------------------------------------------------
+
+class ProgressMonitor:
+    def __init__(self, tag=None, num_items=None, flush_interval=1000, verbose=False, progress_fn=None, pfn_lo=0, pfn_hi=1000, pfn_total=1000):
+        self.tag = tag
+        self.num_items = num_items
+        self.verbose = verbose
+        self.flush_interval = flush_interval
+        self.progress_fn = progress_fn
+        self.pfn_lo = pfn_lo
+        self.pfn_hi = pfn_hi
+        self.pfn_total = pfn_total
+        self.start_time = time.time()
+        self.batch_time = self.start_time
+        self.batch_items = 0
+        if self.progress_fn is not None:
+            self.progress_fn(self.pfn_lo, self.pfn_total)
+
+    def update(self, cur_items):
+        assert (self.num_items is None) or (cur_items <= self.num_items)
+        if (cur_items < self.batch_items + self.flush_interval) and (self.num_items is None or cur_items < self.num_items):
+            return
+        cur_time = time.time()
+        total_time = cur_time - self.start_time
+        time_per_item = (cur_time - self.batch_time) / max(cur_items - self.batch_items, 1)
+        if (self.verbose) and (self.tag is not None):
+            print(f'{self.tag:<19s} items {cur_items:<7d} time {dnnlib.util.format_time(total_time):<12s} ms/item {time_per_item*1e3:.2f}')
+        self.batch_time = cur_time
+        self.batch_items = cur_items
+
+        if (self.progress_fn is not None) and (self.num_items is not None):
+            self.progress_fn(self.pfn_lo + (self.pfn_hi - self.pfn_lo) * (cur_items / self.num_items), self.pfn_total)
+
+    def sub(self, tag=None, num_items=None, flush_interval=1000, rel_lo=0, rel_hi=1):
+        return ProgressMonitor(
+            tag             = tag,
+            num_items       = num_items,
+            flush_interval  = flush_interval,
+            verbose         = self.verbose,
+            progress_fn     = self.progress_fn,
+            pfn_lo          = self.pfn_lo + (self.pfn_hi - self.pfn_lo) * rel_lo,
+            pfn_hi          = self.pfn_lo + (self.pfn_hi - self.pfn_lo) * rel_hi,
+            pfn_total       = self.pfn_total,
+        )
+
+#----------------------------------------------------------------------------
+
+def compute_feature_stats_for_dataset(opts, detector_url, detector_kwargs, rel_lo=0, rel_hi=1, batch_size=64, data_loader_kwargs=None, max_items=None, **stats_kwargs):
+    dataset = dnnlib.util.construct_class_by_name(**opts.dataset_kwargs)
+    if data_loader_kwargs is None:
+        data_loader_kwargs = dict(pin_memory=True, num_workers=3, prefetch_factor=2)
+
+    # Try to lookup from cache.
+    cache_file = None
+    if opts.cache:
+        # Choose cache file name.
+        args = dict(dataset_kwargs=opts.dataset_kwargs, detector_url=detector_url, detector_kwargs=detector_kwargs, stats_kwargs=stats_kwargs)
+        md5 = hashlib.md5(repr(sorted(args.items())).encode('utf-8'))
+        cache_tag = f'{dataset.name}-{get_feature_detector_name(detector_url)}-{md5.hexdigest()}'
+        cache_file = dnnlib.make_cache_dir_path('gan-metrics', cache_tag + '.pkl')
+
+        # Check if the file exists (all processes must agree).
+        flag = os.path.isfile(cache_file) if opts.rank == 0 else False
+        if opts.num_gpus > 1:
+            flag = torch.as_tensor(flag, dtype=torch.float32, device=opts.device)
+            torch.distributed.broadcast(tensor=flag, src=0)
+            flag = (float(flag.cpu()) != 0)
+
+        # Load.
+        if flag:
+            return FeatureStats.load(cache_file)
+
+    # Initialize.
+    num_items = len(dataset)
+    if max_items is not None:
+        num_items = min(num_items, max_items)
+    stats = FeatureStats(max_items=num_items, **stats_kwargs)
+    progress = opts.progress.sub(tag='dataset features', num_items=num_items, rel_lo=rel_lo, rel_hi=rel_hi)
+    detector = get_feature_detector(url=detector_url, device=opts.device, num_gpus=opts.num_gpus, rank=opts.rank, verbose=progress.verbose)
+
+    # Main loop.
+    item_subset = [(i * opts.num_gpus + opts.rank) % num_items for i in range((num_items - 1) // opts.num_gpus + 1)]
+    for images, _labels, _indices in torch.utils.data.DataLoader(dataset=dataset, sampler=item_subset, batch_size=batch_size, **data_loader_kwargs):
+        if images.shape[1] == 1:
+            images = images.repeat([1, 3, 1, 1])
+        features = detector(images.to(opts.device), **detector_kwargs)
+        stats.append_torch(features, num_gpus=opts.num_gpus, rank=opts.rank)
+        progress.update(stats.num_items)
+
+    # Save to cache.
+    if cache_file is not None and opts.rank == 0:
+        os.makedirs(os.path.dirname(cache_file), exist_ok=True)
+        temp_file = cache_file + '.' + uuid.uuid4().hex
+        stats.save(temp_file)
+        os.replace(temp_file, cache_file) # atomic
+    return stats
+
+#----------------------------------------------------------------------------
+
+def compute_feature_stats_for_generator(opts, detector_url, detector_kwargs, rel_lo=0, rel_hi=1, batch_size=64, batch_gen=None, jit=False, **stats_kwargs):
+    if batch_gen is None:
+        batch_gen = min(batch_size, 4)
+    assert batch_size % batch_gen == 0
+
+    # Setup generator and load labels.
+    G = copy.deepcopy(opts.G).eval().requires_grad_(False).to(opts.device)
+    dataset = dnnlib.util.construct_class_by_name(**opts.dataset_kwargs)
+
+    # HACK:
+    # other_data = "/checkpoint/jgu/space/gan/ffhq/giraffe_results/gen_images"
+    # other_data = "/checkpoint/jgu/space/gan/cars/gen_images_380000"
+    # other_data = "/private/home/jgu/work/pi-GAN/Baselines/FFHQEvalOutput2"
+    # other_data = "/private/home/jgu/work/pi-GAN/Baselines/AFHQEvalOutput"
+    # other_data = sorted(glob.glob(f'{other_data}/*.jpg'))
+    # other_data = '/private/home/jgu/work/giraffe/out/afhq256/fid_images.npy'
+    # other_images = np.load(other_data)
+    #　from fairseq import pdb;pdb.set_trace()
+    # print(f'other data size = {len(other_data)}')
+    other_data = None
+
+    # Image generation func.
+    def run_generator(z, c):
+        # from fairseq import pdb;pdb.set_trace()
+        if hasattr(G, 'get_final_output'):
+            img = G.get_final_output(z=z, c=c, **opts.G_kwargs)
+        else:
+            img = G(z=z, c=c, **opts.G_kwargs)
+        img = (img * 127.5 + 128).clamp(0, 255).to(torch.uint8)
+        return img
+
+    # JIT.
+    if jit:
+        z = torch.zeros([batch_gen, G.z_dim], device=opts.device)
+        c = torch.zeros([batch_gen, G.c_dim], device=opts.device)
+        run_generator = torch.jit.trace(run_generator, [z, c], check_trace=False)
+
+    # Initialize.
+    stats = FeatureStats(**stats_kwargs)
+    assert stats.max_items is not None
+    progress = opts.progress.sub(tag='generator features', num_items=stats.max_items, rel_lo=rel_lo, rel_hi=rel_hi)
+    detector = get_feature_detector(url=detector_url, device=opts.device, num_gpus=opts.num_gpus, rank=opts.rank, verbose=progress.verbose)
+
+    # Main loop.
+    till_now = 0
+    while not stats.is_full():
+        images = []
+        if other_data is None:
+            for _i in range(batch_size // batch_gen):
+                z = torch.randn([batch_gen, G.z_dim], device=opts.device)
+                c = [dataset.get_label(np.random.randint(len(dataset))) for _i in range(batch_gen)]
+                c = torch.from_numpy(np.stack(c)).pin_memory().to(opts.device)
+                img = run_generator(z, c)
+                images.append(img)
+            images = torch.cat(images)
+        else:
+            batch_idxs = [((till_now + i) * opts.num_gpus + opts.rank) % len(other_images) for i in range(batch_size)]
+            import imageio
+            till_now += batch_size
+            images = other_images[batch_idxs]
+            images = torch.from_numpy(images).to(opts.device)
+            # images = np.stack([imageio.imread(other_data[i % len(other_data)]) for i in batch_idxs], axis=0)
+            # images = torch.from_numpy(images).to(opts.device).permute(0,3,1,2)
+            
+        if images.shape[1] == 1:
+            images = images.repeat([1, 3, 1, 1])
+        features = detector(images, **detector_kwargs)
+        stats.append_torch(features, num_gpus=opts.num_gpus, rank=opts.rank)
+        progress.update(stats.num_items)
+    return stats
+
+#----------------------------------------------------------------------------

metrics/perceptual_path_length.py

@@ -0,0 +1,131 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+"""Perceptual Path Length (PPL) from the paper "A Style-Based Generator
+Architecture for Generative Adversarial Networks". Matches the original
+implementation by Karras et al. at
+https://github.com/NVlabs/stylegan/blob/master/metrics/perceptual_path_length.py"""
+
+import copy
+import numpy as np
+import torch
+import dnnlib
+from . import metric_utils
+
+#----------------------------------------------------------------------------
+
+# Spherical interpolation of a batch of vectors.
+def slerp(a, b, t):
+    a = a / a.norm(dim=-1, keepdim=True)
+    b = b / b.norm(dim=-1, keepdim=True)
+    d = (a * b).sum(dim=-1, keepdim=True)
+    p = t * torch.acos(d)
+    c = b - d * a
+    c = c / c.norm(dim=-1, keepdim=True)
+    d = a * torch.cos(p) + c * torch.sin(p)
+    d = d / d.norm(dim=-1, keepdim=True)
+    return d
+
+#----------------------------------------------------------------------------
+
+class PPLSampler(torch.nn.Module):
+    def __init__(self, G, G_kwargs, epsilon, space, sampling, crop, vgg16):
+        assert space in ['z', 'w']
+        assert sampling in ['full', 'end']
+        super().__init__()
+        self.G = copy.deepcopy(G)
+        self.G_kwargs = G_kwargs
+        self.epsilon = epsilon
+        self.space = space
+        self.sampling = sampling
+        self.crop = crop
+        self.vgg16 = copy.deepcopy(vgg16)
+
+    def forward(self, c):
+        # Generate random latents and interpolation t-values.
+        t = torch.rand([c.shape[0]], device=c.device) * (1 if self.sampling == 'full' else 0)
+        z0, z1 = torch.randn([c.shape[0] * 2, self.G.z_dim], device=c.device).chunk(2)
+
+        # Interpolate in W or Z.
+        if self.space == 'w':
+            w0, w1 = self.G.mapping(z=torch.cat([z0,z1]), c=torch.cat([c,c])).chunk(2)
+            wt0 = w0.lerp(w1, t.unsqueeze(1).unsqueeze(2))
+            wt1 = w0.lerp(w1, t.unsqueeze(1).unsqueeze(2) + self.epsilon)
+        else: # space == 'z'
+            zt0 = slerp(z0, z1, t.unsqueeze(1))
+            zt1 = slerp(z0, z1, t.unsqueeze(1) + self.epsilon)
+            wt0, wt1 = self.G.mapping(z=torch.cat([zt0,zt1]), c=torch.cat([c,c])).chunk(2)
+
+        # Randomize noise buffers.
+        for name, buf in self.G.named_buffers():
+            if name.endswith('.noise_const'):
+                buf.copy_(torch.randn_like(buf))
+
+        # Generate images.
+        img = self.G.synthesis(ws=torch.cat([wt0,wt1]), noise_mode='const', force_fp32=True, **self.G_kwargs)
+
+        # Center crop.
+        if self.crop:
+            assert img.shape[2] == img.shape[3]
+            c = img.shape[2] // 8
+            img = img[:, :, c*3 : c*7, c*2 : c*6]
+
+        # Downsample to 256x256.
+        factor = self.G.img_resolution // 256
+        if factor > 1:
+            img = img.reshape([-1, img.shape[1], img.shape[2] // factor, factor, img.shape[3] // factor, factor]).mean([3, 5])
+
+        # Scale dynamic range from [-1,1] to [0,255].
+        img = (img + 1) * (255 / 2)
+        if self.G.img_channels == 1:
+            img = img.repeat([1, 3, 1, 1])
+
+        # Evaluate differential LPIPS.
+        lpips_t0, lpips_t1 = self.vgg16(img, resize_images=False, return_lpips=True).chunk(2)
+        dist = (lpips_t0 - lpips_t1).square().sum(1) / self.epsilon ** 2
+        return dist
+
+#----------------------------------------------------------------------------
+
+def compute_ppl(opts, num_samples, epsilon, space, sampling, crop, batch_size, jit=False):
+    dataset = dnnlib.util.construct_class_by_name(**opts.dataset_kwargs)
+    vgg16_url = 'https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/metrics/vgg16.pt'
+    vgg16 = metric_utils.get_feature_detector(vgg16_url, num_gpus=opts.num_gpus, rank=opts.rank, verbose=opts.progress.verbose)
+
+    # Setup sampler.
+    sampler = PPLSampler(G=opts.G, G_kwargs=opts.G_kwargs, epsilon=epsilon, space=space, sampling=sampling, crop=crop, vgg16=vgg16)
+    sampler.eval().requires_grad_(False).to(opts.device)
+    if jit:
+        c = torch.zeros([batch_size, opts.G.c_dim], device=opts.device)
+        sampler = torch.jit.trace(sampler, [c], check_trace=False)
+
+    # Sampling loop.
+    dist = []
+    progress = opts.progress.sub(tag='ppl sampling', num_items=num_samples)
+    for batch_start in range(0, num_samples, batch_size * opts.num_gpus):
+        progress.update(batch_start)
+        c = [dataset.get_label(np.random.randint(len(dataset))) for _i in range(batch_size)]
+        c = torch.from_numpy(np.stack(c)).pin_memory().to(opts.device)
+        x = sampler(c)
+        for src in range(opts.num_gpus):
+            y = x.clone()
+            if opts.num_gpus > 1:
+                torch.distributed.broadcast(y, src=src)
+            dist.append(y)
+    progress.update(num_samples)
+
+    # Compute PPL.
+    if opts.rank != 0:
+        return float('nan')
+    dist = torch.cat(dist)[:num_samples].cpu().numpy()
+    lo = np.percentile(dist, 1, interpolation='lower')
+    hi = np.percentile(dist, 99, interpolation='higher')
+    ppl = np.extract(np.logical_and(dist >= lo, dist <= hi), dist).mean()
+    return float(ppl)
+
+#----------------------------------------------------------------------------

metrics/precision_recall.py

@@ -0,0 +1,62 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+"""Precision/Recall (PR) from the paper "Improved Precision and Recall
+Metric for Assessing Generative Models". Matches the original implementation
+by Kynkaanniemi et al. at
+https://github.com/kynkaat/improved-precision-and-recall-metric/blob/master/precision_recall.py"""
+
+import torch
+from . import metric_utils
+
+#----------------------------------------------------------------------------
+
+def compute_distances(row_features, col_features, num_gpus, rank, col_batch_size):
+    assert 0 <= rank < num_gpus
+    num_cols = col_features.shape[0]
+    num_batches = ((num_cols - 1) // col_batch_size // num_gpus + 1) * num_gpus
+    col_batches = torch.nn.functional.pad(col_features, [0, 0, 0, -num_cols % num_batches]).chunk(num_batches)
+    dist_batches = []
+    for col_batch in col_batches[rank :: num_gpus]:
+        dist_batch = torch.cdist(row_features.unsqueeze(0), col_batch.unsqueeze(0))[0]
+        for src in range(num_gpus):
+            dist_broadcast = dist_batch.clone()
+            if num_gpus > 1:
+                torch.distributed.broadcast(dist_broadcast, src=src)
+            dist_batches.append(dist_broadcast.cpu() if rank == 0 else None)
+    return torch.cat(dist_batches, dim=1)[:, :num_cols] if rank == 0 else None
+
+#----------------------------------------------------------------------------
+
+def compute_pr(opts, max_real, num_gen, nhood_size, row_batch_size, col_batch_size):
+    detector_url = 'https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/metrics/vgg16.pt'
+    detector_kwargs = dict(return_features=True)
+
+    real_features = metric_utils.compute_feature_stats_for_dataset(
+        opts=opts, detector_url=detector_url, detector_kwargs=detector_kwargs,
+        rel_lo=0, rel_hi=0, capture_all=True, max_items=max_real).get_all_torch().to(torch.float16).to(opts.device)
+
+    gen_features = metric_utils.compute_feature_stats_for_generator(
+        opts=opts, detector_url=detector_url, detector_kwargs=detector_kwargs,
+        rel_lo=0, rel_hi=1, capture_all=True, max_items=num_gen).get_all_torch().to(torch.float16).to(opts.device)
+
+    results = dict()
+    for name, manifold, probes in [('precision', real_features, gen_features), ('recall', gen_features, real_features)]:
+        kth = []
+        for manifold_batch in manifold.split(row_batch_size):
+            dist = compute_distances(row_features=manifold_batch, col_features=manifold, num_gpus=opts.num_gpus, rank=opts.rank, col_batch_size=col_batch_size)
+            kth.append(dist.to(torch.float32).kthvalue(nhood_size + 1).values.to(torch.float16) if opts.rank == 0 else None)
+        kth = torch.cat(kth) if opts.rank == 0 else None
+        pred = []
+        for probes_batch in probes.split(row_batch_size):
+            dist = compute_distances(row_features=probes_batch, col_features=manifold, num_gpus=opts.num_gpus, rank=opts.rank, col_batch_size=col_batch_size)
+            pred.append((dist <= kth).any(dim=1) if opts.rank == 0 else None)
+        results[name] = float(torch.cat(pred).to(torch.float32).mean() if opts.rank == 0 else 'nan')
+    return results['precision'], results['recall']
+
+#----------------------------------------------------------------------------

renderer.py

@@ -0,0 +1,322 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+
+"""Wrap the generator to render a sequence of images"""
+import torch
+import torch.nn.functional as F
+import numpy as np
+from torch import random
+import tqdm
+import copy
+import trimesh
+
+
+class Renderer(object):
+
+    def __init__(self, generator, discriminator=None, program=None):
+        self.generator = generator
+        self.discriminator = discriminator
+        self.sample_tmp = 0.65
+        self.program = program
+        self.seed = 0
+
+        if (program is not None) and (len(program.split(':')) == 2):
+            from training.dataset import ImageFolderDataset
+            self.image_data = ImageFolderDataset(program.split(':')[1])
+            self.program = program.split(':')[0]
+        else:
+            self.image_data = None
+
+    def set_random_seed(self, seed):
+        self.seed = seed
+        torch.manual_seed(seed)
+        np.random.seed(seed)
+
+    def __call__(self, *args, **kwargs):
+        self.generator.eval()  # eval mode...
+
+        if self.program is None:
+            if hasattr(self.generator, 'get_final_output'):
+                return self.generator.get_final_output(*args, **kwargs)
+            return self.generator(*args, **kwargs)
+        
+        if self.image_data is not None:
+            batch_size = 1
+            indices = (np.random.rand(batch_size) * len(self.image_data)).tolist()
+            rimages = np.stack([self.image_data._load_raw_image(int(i)) for i in indices], 0)
+            rimages = torch.from_numpy(rimages).float().to(kwargs['z'].device) / 127.5 - 1
+            kwargs['img'] = rimages
+        
+        outputs = getattr(self, f"render_{self.program}")(*args, **kwargs)
+        
+        if self.image_data is not None:
+            imgs = outputs if not isinstance(outputs, tuple) else outputs[0]
+            size = imgs[0].size(-1)
+            rimg = F.interpolate(rimages, (size, size), mode='bicubic', align_corners=False)
+            imgs = [torch.cat([img, rimg], 0) for img in imgs]
+            outputs = imgs if not isinstance(outputs, tuple) else (imgs, outputs[1])
+        return outputs
+
+    def get_additional_params(self, ws, t=0):
+        gen = self.generator.synthesis
+        batch_size = ws.size(0)
+
+        kwargs = {}
+        if not hasattr(gen, 'get_latent_codes'):
+            return kwargs
+
+        s_val, t_val, r_val = [[0, 0, 0]], [[0.5, 0.5, 0.5]], [0.]
+        # kwargs["transformations"] = gen.get_transformations(batch_size=batch_size, mode=[s_val, t_val, r_val], device=ws.device)
+        # kwargs["bg_rotation"] = gen.get_bg_rotation(batch_size, device=ws.device)
+        # kwargs["light_dir"] = gen.get_light_dir(batch_size, device=ws.device)
+        kwargs["latent_codes"] = gen.get_latent_codes(batch_size, tmp=self.sample_tmp, device=ws.device)
+        kwargs["camera_matrices"] = self.get_camera_traj(t, ws.size(0), device=ws.device)
+        return kwargs
+
+    def get_camera_traj(self, t, batch_size=1, traj_type='pigan', device='cpu'):
+        gen = self.generator.synthesis
+        if traj_type == 'pigan':
+            range_u, range_v = gen.C.range_u, gen.C.range_v
+            pitch = 0.2 * np.cos(t * 2 * np.pi) + np.pi/2
+            yaw = 0.4 * np.sin(t * 2 * np.pi)
+            u = (yaw - range_u[0]) / (range_u[1] - range_u[0])
+            v = (pitch - range_v[0]) / (range_v[1] - range_v[0])
+            cam = gen.get_camera(batch_size=batch_size, mode=[u, v, 0.5], device=device)
+        else:
+            raise NotImplementedError
+        return cam
+   
+    def render_rotation_camera(self, *args, **kwargs):
+        batch_size, n_steps = 2, kwargs["n_steps"]
+        gen = self.generator.synthesis
+
+        if 'img' not in kwargs:
+            ws = self.generator.mapping(*args, **kwargs)
+        else:
+            ws, _ = self.generator.encoder(kwargs['img'])
+        # ws = ws.repeat(batch_size, 1, 1)
+
+        # kwargs["not_render_background"] = True
+        if hasattr(gen, 'get_latent_codes'):
+            kwargs["latent_codes"] = gen.get_latent_codes(batch_size, tmp=self.sample_tmp, device=ws.device)
+            kwargs.pop('img', None) 
+
+        out = []
+        cameras = []
+        relatve_range_u = kwargs['relative_range_u']
+        u_samples = np.linspace(relatve_range_u[0], relatve_range_u[1], n_steps)
+        for step in tqdm.tqdm(range(n_steps)):
+            # Set Camera
+            u = u_samples[step]
+            kwargs["camera_matrices"] = gen.get_camera(batch_size=batch_size, mode=[u, 0.5, 0.5], device=ws.device)
+            cameras.append(gen.get_camera(batch_size=batch_size, mode=[u, 0.5, 0.5], device=ws.device))
+            with torch.no_grad():
+                out_i = gen(ws, **kwargs)
+                if isinstance(out_i, dict):
+                    out_i = out_i['img']
+            out.append(out_i)
+
+        if 'return_cameras' in kwargs and kwargs["return_cameras"]:
+            return out, cameras
+        else:
+            return out
+
+    def render_rotation_camera3(self, styles=None, *args, **kwargs): 
+        gen = self.generator.synthesis
+        n_steps = 36  # 120
+
+        if styles is None:
+            batch_size = 2
+            if 'img' not in kwargs:
+                ws = self.generator.mapping(*args, **kwargs)
+            else:
+                ws = self.generator.encoder(kwargs['img'])['ws']
+            # ws = ws.repeat(batch_size, 1, 1)
+        else:
+            ws = styles
+            batch_size = ws.size(0)
+
+        # kwargs["not_render_background"] = True
+        # Get Random codes and bg rotation
+        self.sample_tmp = 0.72
+        if hasattr(gen, 'get_latent_codes'):
+            kwargs["latent_codes"] = gen.get_latent_codes(batch_size, tmp=self.sample_tmp, device=ws.device)
+            kwargs.pop('img', None) 
+
+        # if getattr(gen, "use_noise", False):
+        #     from dnnlib.geometry import extract_geometry
+        #     kwargs['meshes'] = {}
+        #     low_res, high_res = gen.resolution_vol, gen.img_resolution
+        #     res = low_res * 2
+        #     while res <= high_res:
+        #         kwargs['meshes'][res] = [trimesh.Trimesh(*extract_geometry(gen, ws, resolution=res, threshold=30.))]
+        #         kwargs['meshes'][res] += [
+        #             torch.randn(len(kwargs['meshes'][res][0].vertices), 
+        #                 2, device=ws.device)[kwargs['meshes'][res][0].faces]]
+        #         res = res * 2
+        # if getattr(gen, "use_noise", False):
+        #     kwargs['voxel_noise'] = gen.get_voxel_field(styles=ws, n_vols=2048, return_noise=True, sphere_noise=True)
+        # if getattr(gen, "use_voxel_noise", False):
+        #     kwargs['voxel_noise'] = gen.get_voxel_field(styles=ws, n_vols=128, return_noise=True)
+        kwargs['noise_mode'] = 'const'
+        
+        out = []
+        tspace = np.linspace(0, 1, n_steps)
+        range_u, range_v = gen.C.range_u, gen.C.range_v
+        
+        for step in tqdm.tqdm(range(n_steps)):
+            t = tspace[step]
+            pitch = 0.2 * np.cos(t * 2 * np.pi) + np.pi/2
+            yaw = 0.4 * np.sin(t * 2 * np.pi)
+            u = (yaw - range_u[0]) / (range_u[1] - range_u[0])
+            v = (pitch - range_v[0]) / (range_v[1] - range_v[0])
+            
+            kwargs["camera_matrices"] = gen.get_camera(
+                batch_size=batch_size, mode=[u, v, t], device=ws.device)
+            
+            with torch.no_grad():
+                out_i = gen(ws, **kwargs)
+                if isinstance(out_i, dict):
+                    out_i = out_i['img'] 
+            out.append(out_i)
+        return out
+
+    def render_rotation_both(self, *args, **kwargs): 
+        gen = self.generator.synthesis
+        batch_size, n_steps = 1, 36 
+        if 'img' not in kwargs:
+            ws = self.generator.mapping(*args, **kwargs)
+        else:
+            ws, _ = self.generator.encoder(kwargs['img'])
+        ws = ws.repeat(batch_size, 1, 1)
+
+        # kwargs["not_render_background"] = True
+        # Get Random codes and bg rotation
+        kwargs["latent_codes"] = gen.get_latent_codes(batch_size, tmp=self.sample_tmp, device=ws.device)
+        kwargs.pop('img', None) 
+
+        out = []
+        tspace = np.linspace(0, 1, n_steps)
+        range_u, range_v = gen.C.range_u, gen.C.range_v
+        
+        for step in tqdm.tqdm(range(n_steps)):
+            t = tspace[step]
+            pitch = 0.2 * np.cos(t * 2 * np.pi) + np.pi/2
+            yaw = 0.4 * np.sin(t * 2 * np.pi)
+            u = (yaw - range_u[0]) / (range_u[1] - range_u[0])
+            v = (pitch - range_v[0]) / (range_v[1] - range_v[0])
+
+            kwargs["camera_matrices"] = gen.get_camera(
+                batch_size=batch_size, mode=[u, v, 0.5], device=ws.device)
+            
+            with torch.no_grad():
+                out_i = gen(ws, **kwargs)
+                if isinstance(out_i, dict):
+                    out_i = out_i['img']  
+
+                kwargs_n = copy.deepcopy(kwargs)
+                kwargs_n.update({'render_option': 'early,no_background,up64,depth,normal'})               
+                out_n = gen(ws, **kwargs_n)
+                out_n = F.interpolate(out_n, 
+                    size=(out_i.size(-1), out_i.size(-1)), 
+                    mode='bicubic', align_corners=True)
+                out_i = torch.cat([out_i, out_n], 0)
+            out.append(out_i)
+        return out
+
+    def render_rotation_grid(self, styles=None, return_cameras=False, *args, **kwargs):
+        gen = self.generator.synthesis
+        if styles is None:
+            batch_size = 1
+            ws = self.generator.mapping(*args, **kwargs)
+            ws = ws.repeat(batch_size, 1, 1)
+        else:
+            ws = styles
+            batch_size = ws.size(0)
+
+        kwargs["latent_codes"] = gen.get_latent_codes(batch_size, tmp=self.sample_tmp, device=ws.device)
+        kwargs.pop('img', None) 
+
+        if getattr(gen, "use_voxel_noise", False):
+            kwargs['voxel_noise'] = gen.get_voxel_field(styles=ws, n_vols=128, return_noise=True)
+
+        out = []
+        cameras = []
+        range_u, range_v = gen.C.range_u, gen.C.range_v
+
+        a_steps, b_steps = 6, 3
+        aspace = np.linspace(-0.4, 0.4, a_steps)
+        bspace = np.linspace(-0.2, 0.2, b_steps) * -1
+        for b in tqdm.tqdm(range(b_steps)):
+            for a in range(a_steps):
+                t_a = aspace[a]
+                t_b = bspace[b]
+                camera_mat = gen.camera_matrix.repeat(batch_size, 1, 1).to(ws.device)
+                loc_x = np.cos(t_b) * np.cos(t_a)
+                loc_y = np.cos(t_b) * np.sin(t_a)
+                loc_z = np.sin(t_b)
+                loc = torch.tensor([[loc_x, loc_y, loc_z]], dtype=torch.float32).to(ws.device)
+                from dnnlib.camera import look_at
+                R = look_at(loc)
+                RT = torch.eye(4).reshape(1, 4, 4).repeat(batch_size, 1, 1)
+                RT[:, :3, :3] = R
+                RT[:, :3, -1] = loc
+
+                world_mat = RT.to(ws.device)
+                #kwargs["camera_matrices"] = gen.get_camera(
+                #     batch_size=batch_size, mode=[u, v, 0.5], device=ws.device)
+                kwargs["camera_matrices"] = (camera_mat, world_mat, "random", None)
+
+                with torch.no_grad():
+                    out_i = gen(ws, **kwargs)
+                    if isinstance(out_i, dict):
+                        out_i = out_i['img']
+
+                    # kwargs_n = copy.deepcopy(kwargs)
+                    # kwargs_n.update({'render_option': 'early,no_background,up64,depth,normal'})
+                    # out_n = gen(ws, **kwargs_n)
+                    # out_n = F.interpolate(out_n,
+                    #                       size=(out_i.size(-1), out_i.size(-1)),
+                    #                       mode='bicubic', align_corners=True)
+                    # out_i = torch.cat([out_i, out_n], 0)
+                out.append(out_i)
+
+        if return_cameras:
+            return out, cameras
+        else:
+            return out
+
+    def render_rotation_camera_grid(self, *args, **kwargs): 
+        batch_size, n_steps = 1, 60
+        gen = self.generator.synthesis
+        bbox_generator = self.generator.synthesis.boundingbox_generator
+        
+        ws = self.generator.mapping(*args, **kwargs)
+        ws = ws.repeat(batch_size, 1, 1)
+
+        # Get Random codes and bg rotation
+        kwargs["latent_codes"] = gen.get_latent_codes(batch_size, tmp=self.sample_tmp, device=ws.device)
+        del kwargs['render_option']
+
+        out = []
+        for v in [0.15, 0.5, 1.05]:
+            for step in tqdm.tqdm(range(n_steps)):
+                # Set Camera
+                u = step * 1.0 / (n_steps - 1) - 1.0 
+                kwargs["camera_matrices"] = gen.get_camera(batch_size=batch_size, mode=[u, v, 0.5], device=ws.device)
+                with torch.no_grad():
+                    out_i = gen(ws, render_option=None, **kwargs)
+                    if isinstance(out_i, dict):
+                        out_i = out_i['img']
+                    # option_n = 'early,no_background,up64,depth,direct_depth'
+                    # option_n = 'early,up128,no_background,depth,normal'                
+                    # out_n = gen(ws, render_option=option_n, **kwargs)
+                    # out_n = F.interpolate(out_n, 
+                    #     size=(out_i.size(-1), out_i.size(-1)), 
+                    #     mode='bicubic', align_corners=True)
+                    # out_i = torch.cat([out_i, out_n], 0)
+            
+                out.append(out_i)
+
+        # out += out[::-1]
+        return out

requirements.txt

@@ -0,0 +1,29 @@
+albumentations==1.0.3
+click==8.0.3
+clip-by-openai==1.1
+einops==0.3.0
+glfw==2.5.0
+gradio==2.8.13
+imageio==2.9.0
+imgui==1.4.1
+kornia==0.5.10
+lmdb==0.98
+lpips==0.1.4
+matplotlib==3.4.3
+numpy==1.21.2
+hydra-core==1.1
+opencv_python_headless==4.5.1.48
+Pillow==9.0.1
+psutil==5.8.0
+PyMCubes==0.1.2
+PyOpenGL==3.1.6
+pyspng==0.1.0
+requests==2.26.0
+scipy==1.7.1
+submitit==1.1.5
+tensorboardX==2.5
+torch==1.7.1
+torchvision==0.8.2
+tqdm==4.62.2
+trimesh==3.9.8
+imageio-ffmpeg==0.4.5

run_train.py

@@ -0,0 +1,398 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+
+from math import dist
+import sys
+import os
+import click
+import re
+import json
+import glob
+import tempfile
+import torch
+import dnnlib
+import hydra
+
+from datetime import date
+from training import training_loop
+from metrics import metric_main
+from torch_utils import training_stats, custom_ops, distributed_utils
+from torch_utils.distributed_utils import get_init_file, get_shared_folder
+from omegaconf import DictConfig, OmegaConf
+
+#----------------------------------------------------------------------------
+
+class UserError(Exception):
+    pass
+
+#----------------------------------------------------------------------------
+
+def setup_training_loop_kwargs(cfg):
+    args = OmegaConf.create({})
+
+    # ------------------------------------------
+    # General options: gpus, snap, metrics, seed
+    # ------------------------------------------
+    args.rank       = 0
+    args.gpu        = 0
+    args.num_gpus   = torch.cuda.device_count() if cfg.gpus is None else cfg.gpus
+    args.nodes      = cfg.nodes if cfg.nodes is not None else 1
+    args.world_size = 1
+    
+    args.dist_url   = 'env://'
+    args.launcher   = cfg.launcher
+    args.partition  = cfg.partition
+    args.comment    = cfg.comment
+    args.timeout    = 4320 if cfg.timeout is None else cfg.timeout
+    args.job_dir    = ''
+
+    if cfg.snap is None:
+        cfg.snap = 50
+    assert isinstance(cfg.snap, int)
+    if cfg.snap < 1:
+        raise UserError('snap must be at least 1')
+    args.image_snapshot_ticks = cfg.imgsnap
+    args.network_snapshot_ticks = cfg.snap
+    if hasattr(cfg, 'ucp'):
+        args.update_cam_prior_ticks = cfg.ucp
+
+    if cfg.metrics is None:
+        cfg.metrics = ['fid50k_full']
+    cfg.metrics = list(cfg.metrics)
+    if not all(metric_main.is_valid_metric(metric) for metric in cfg.metrics):
+        raise UserError('\n'.join(['metrics can only contain the following values:'] + metric_main.list_valid_metrics()))
+    args.metrics = cfg.metrics
+
+    if cfg.seed is None:
+        cfg.seed = 0
+    assert isinstance(cfg.seed, int)
+    args.random_seed = cfg.seed
+
+    # -----------------------------------
+    # Dataset: data, cond, subset, mirror
+    # ----------------------------------- 
+
+    assert cfg.data is not None
+    assert isinstance(cfg.data, str)
+    args.update({"training_set_kwargs": dict(class_name='training.dataset.ImageFolderDataset', path=cfg.data, resolution=cfg.resolution, use_labels=True, max_size=None, xflip=False)})
+    args.update({"data_loader_kwargs": dict(pin_memory=True, num_workers=3, prefetch_factor=2)})
+    args.generation_with_image = getattr(cfg, 'generate_with_image', False)
+    try:
+        training_set = dnnlib.util.construct_class_by_name(**args.training_set_kwargs) # subclass of training.dataset.Dataset
+        args.training_set_kwargs.resolution = training_set.resolution                  # be explicit about resolution
+        args.training_set_kwargs.use_labels = training_set.has_labels                  # be explicit about labels
+        args.training_set_kwargs.max_size = len(training_set)                          # be explicit about dataset size
+        desc = training_set.name
+        del training_set # conserve memory
+    except IOError as err:
+        raise UserError(f'data: {err}')
+
+    if cfg.cond is None:
+        cfg.cond = False
+    assert isinstance(cfg.cond, bool)
+    if cfg.cond:
+        if not args.training_set_kwargs.use_labels:
+            raise UserError('cond=True requires labels specified in dataset.json')
+        desc += '-cond'
+    else:
+        args.training_set_kwargs.use_labels = False
+
+    if cfg.subset is not None:
+        assert isinstance(cfg.subset, int)
+        if not 1 <= cfg.subset <= args.training_set_kwargs.max_size:
+            raise UserError(f'subset must be between 1 and {args.training_set_kwargs.max_size}')
+        desc += f'-subset{cfg.subset}'
+        if cfg.subset < args.training_set_kwargs.max_size:
+            args.training_set_kwargs.max_size = cfg.subset
+            args.training_set_kwargs.random_seed = args.random_seed
+
+    if cfg.mirror is None:
+        cfg.mirror = False
+    assert isinstance(cfg.mirror, bool)
+    if cfg.mirror:
+        desc += '-mirror'
+        args.training_set_kwargs.xflip = True
+
+    # ------------------------------------
+    # Base config: cfg, model, gamma, kimg, batch
+    # ------------------------------------
+    if cfg.auto:
+        cfg.spec.name = 'auto'
+    desc += f'-{cfg.spec.name}'
+    desc += f'-{cfg.model.name}'
+    if cfg.spec.name == 'auto':
+        res = args.training_set_kwargs.resolution
+        cfg.spec.fmaps = 1 if res >= 512 else 0.5
+        cfg.spec.lrate = 0.002 if res >= 1024 else 0.0025
+        cfg.spec.gamma = 0.0002 * (res ** 2) / cfg.spec.mb # heuristic formula
+        cfg.spec.ema = cfg.spec.mb * 10 / 32
+    
+    if getattr(cfg.spec, 'lrate_disc', None) is None:
+        cfg.spec.lrate_disc = cfg.spec.lrate   # use the same learning rate for discriminator
+
+    # model (generator, discriminator)
+    args.update({"G_kwargs": dict(**cfg.model.G_kwargs)})
+    args.update({"D_kwargs": dict(**cfg.model.D_kwargs)})
+    args.update({"G_opt_kwargs": dict(class_name='torch.optim.Adam', lr=cfg.spec.lrate, betas=[0,0.99], eps=1e-8)})
+    args.update({"D_opt_kwargs": dict(class_name='torch.optim.Adam', lr=cfg.spec.lrate_disc, betas=[0,0.99], eps=1e-8)})
+    args.update({"loss_kwargs": dict(class_name='training.loss.StyleGAN2Loss', r1_gamma=cfg.spec.gamma, **cfg.model.loss_kwargs)})
+    
+    if cfg.spec.name == 'cifar':
+        args.loss_kwargs.pl_weight = 0 # disable path length regularization
+        args.loss_kwargs.style_mixing_prob = 0 # disable style mixing
+        args.D_kwargs.architecture = 'orig' # disable residual skip connections
+
+    # kimg data config
+    args.spec = cfg.spec  # just keep the dict.
+    args.total_kimg = cfg.spec.kimg
+    args.batch_size = cfg.spec.mb
+    args.batch_gpu = cfg.spec.mbstd
+    args.ema_kimg = cfg.spec.ema
+    args.ema_rampup = cfg.spec.ramp
+    
+    # ---------------------------------------------------
+    # Discriminator augmentation: aug, p, target, augpipe
+    # ---------------------------------------------------
+    if cfg.aug is None:
+        cfg.aug = 'ada'
+    else:
+        assert isinstance(cfg.aug, str)
+        desc += f'-{cfg.aug}'
+
+    if cfg.aug == 'ada':
+        args.ada_target = 0.6
+    elif cfg.aug == 'noaug':
+        pass
+    elif cfg.aug == 'fixed':
+        if cfg.p is None:
+            raise UserError(f'--aug={cfg.aug} requires specifying --p')
+    else:
+        raise UserError(f'--aug={cfg.aug} not supported')
+
+    if cfg.p is not None:
+        assert isinstance(cfg.p, float)
+        if cfg.aug != 'fixed':
+            raise UserError('--p can only be specified with --aug=fixed')
+        if not 0 <= cfg.p <= 1:
+            raise UserError('--p must be between 0 and 1')
+        desc += f'-p{cfg.p:g}'
+        args.augment_p = cfg.p
+
+    if cfg.target is not None:
+        assert isinstance(cfg.target, float)
+        if cfg.aug != 'ada':
+            raise UserError('--target can only be specified with --aug=ada')
+        if not 0 <= cfg.target <= 1:
+            raise UserError('--target must be between 0 and 1')
+        desc += f'-target{cfg.target:g}'
+        args.ada_target = cfg.target
+
+    assert cfg.augpipe is None or isinstance(cfg.augpipe, str)
+    if cfg.augpipe is None:
+        cfg.augpipe = 'bgc'
+    else:
+        if cfg.aug == 'noaug':
+            raise UserError('--augpipe cannot be specified with --aug=noaug')
+        desc += f'-{cfg.augpipe}'
+
+    augpipe_specs = {
+        'blit':   dict(xflip=1, rotate90=1, xint=1),
+        'geom':   dict(scale=1, rotate=1, aniso=1, xfrac=1),
+        'color':  dict(brightness=1, contrast=1, lumaflip=1, hue=1, saturation=1),
+        'filter': dict(imgfilter=1),
+        'noise':  dict(noise=1),
+        'cutout': dict(cutout=1),
+        'bgc0':   dict(xint=1, scale=1, aniso=1, xfrac=1, brightness=1, contrast=1, lumaflip=1, hue=1, saturation=1),
+        'bg':     dict(xflip=1, rotate90=1, xint=1, scale=1, rotate=1, aniso=1, xfrac=1),
+        'bgc':    dict(xflip=1, rotate90=1, xint=1, scale=1, rotate=1, aniso=1, xfrac=1, brightness=1, contrast=1, lumaflip=1, hue=1, saturation=1),
+        'bgcf':   dict(xflip=1, rotate90=1, xint=1, scale=1, rotate=1, aniso=1, xfrac=1, brightness=1, contrast=1, lumaflip=1, hue=1, saturation=1, imgfilter=1),
+        'bgcfn':  dict(xflip=1, rotate90=1, xint=1, scale=1, rotate=1, aniso=1, xfrac=1, brightness=1, contrast=1, lumaflip=1, hue=1, saturation=1, imgfilter=1, noise=1),
+        'bgcfnc': dict(xflip=1, rotate90=1, xint=1, scale=1, rotate=1, aniso=1, xfrac=1, brightness=1, contrast=1, lumaflip=1, hue=1, saturation=1, imgfilter=1, noise=1, cutout=1),
+    }
+    assert cfg.augpipe in augpipe_specs
+    if cfg.aug != 'noaug':
+        args.update({"augment_kwargs": dict(class_name='training.augment.AugmentPipe', **augpipe_specs[cfg.augpipe])})
+
+    # ----------------------------------
+    # Transfer learning: resume, freezed
+    # ----------------------------------
+
+    resume_specs = {
+        'ffhq256':     'https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/transfer-learning-source-nets/ffhq-res256-mirror-paper256-noaug.pkl',
+        'ffhq512':     'https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/transfer-learning-source-nets/ffhq-res512-mirror-stylegan2-noaug.pkl',
+        'ffhq1024':    'https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/transfer-learning-source-nets/ffhq-res1024-mirror-stylegan2-noaug.pkl',
+        'celebahq256': 'https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/transfer-learning-source-nets/celebahq-res256-mirror-paper256-kimg100000-ada-target0.5.pkl',
+        'lsundog256':  'https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/transfer-learning-source-nets/lsundog-res256-paper256-kimg100000-noaug.pkl',
+    }
+
+    assert cfg.resume is None or isinstance(cfg.resume, str)
+    if cfg.resume is None:
+        cfg.resume = 'noresume'
+    elif cfg.resume == 'noresume':
+        desc += '-noresume'
+    elif cfg.resume in resume_specs:
+        desc += f'-resume{cfg.resume}'
+        args.resume_pkl = resume_specs[cfg.resume] # predefined url
+    else:
+        desc += '-resumecustom'
+        args.resume_pkl = cfg.resume # custom path or url
+
+    if cfg.resume != 'noresume':
+        args.ada_kimg = 100 # make ADA react faster at the beginning
+        args.ema_rampup = None # disable EMA rampup
+
+    if cfg.freezed is not None:
+        assert isinstance(cfg.freezed, int)
+        if not cfg.freezed >= 0:
+            raise UserError('--freezed must be non-negative')
+        desc += f'-freezed{cfg.freezed:d}'
+        args.D_kwargs.block_kwargs.freeze_layers = cfg.freezed
+
+    # -------------------------------------------------
+    # Performance options: fp32, nhwc, nobench, workers
+    # -------------------------------------------------
+    args.num_fp16_res = cfg.num_fp16_res
+    if cfg.fp32 is None:
+        cfg.fp32 = False
+    assert isinstance(cfg.fp32, bool)
+    if cfg.fp32:
+        args.G_kwargs.synthesis_kwargs.num_fp16_res = args.D_kwargs.num_fp16_res = 0
+        args.G_kwargs.synthesis_kwargs.conv_clamp = args.D_kwargs.conv_clamp = None
+
+    if cfg.nhwc is None:
+        cfg.nhwc = False
+    assert isinstance(cfg.nhwc, bool)
+    if cfg.nhwc:
+        args.G_kwargs.synthesis_kwargs.fp16_channels_last = args.D_kwargs.block_kwargs.fp16_channels_last = True
+
+    if cfg.nobench is None:
+        cfg.nobench = False
+    assert isinstance(cfg.nobench, bool)
+    if cfg.nobench:
+        args.cudnn_benchmark = False
+
+    if cfg.allow_tf32 is None:
+        cfg.allow_tf32 = False
+    assert isinstance(cfg.allow_tf32, bool)
+    args.allow_tf32 = cfg.allow_tf32
+
+    if cfg.workers is not None:
+        assert isinstance(cfg.workers, int)
+        if not cfg.workers >= 1:
+            raise UserError('--workers must be at least 1')
+        args.data_loader_kwargs.num_workers = cfg.workers
+
+    args.debug = cfg.debug
+    if getattr(cfg, "prefix", None) is not None:
+        desc = cfg.prefix + '-' + desc
+    return desc, args
+
+#----------------------------------------------------------------------------
+
+def subprocess_fn(rank, args):
+    if not args.debug:
+        dnnlib.util.Logger(file_name=os.path.join(args.run_dir, 'log.txt'), file_mode='a', should_flush=True)
+
+    # Init torch.distributed.
+    distributed_utils.init_distributed_mode(rank, args)
+    if args.rank != 0:
+        custom_ops.verbosity = 'none'
+    
+    # Execute training loop.
+    training_loop.training_loop(**args)
+    
+#----------------------------------------------------------------------------
+
+class CommaSeparatedList(click.ParamType):
+    name = 'list'
+
+    def convert(self, value, param, ctx):
+        _ = param, ctx
+        if value is None or value.lower() == 'none' or value == '':
+            return []
+        return value.split(',')
+
+
+@hydra.main(config_path="conf", config_name="config")
+def main(cfg: DictConfig):
+    
+    outdir = cfg.outdir
+
+    # Setup training options
+    run_desc, args = setup_training_loop_kwargs(cfg)
+    
+    # Pick output directory.
+    prev_run_dirs = []
+    if os.path.isdir(outdir):
+        prev_run_dirs = [x for x in os.listdir(outdir) if os.path.isdir(os.path.join(outdir, x))]
+    
+    if cfg.resume_run is None:
+        prev_run_ids = [re.match(r'^\d+', x) for x in prev_run_dirs]
+        prev_run_ids = [int(x.group()) for x in prev_run_ids if x is not None]
+        cur_run_id = max(prev_run_ids, default=-1) + 1
+    else:
+        cur_run_id = cfg.resume_run
+        
+    args.run_dir = os.path.join(outdir, f'{cur_run_id:05d}-{run_desc}')
+    print(outdir, args.run_dir)
+
+    if cfg.resume_run is not None:
+        pkls = sorted(glob.glob(args.run_dir + '/network*.pkl'))
+        if len(pkls) > 0:
+            args.resume_pkl = pkls[-1]
+            args.resume_start = int(args.resume_pkl.split('-')[-1][:-4]) * 1000
+        else:
+            args.resume_start = 0
+
+    # Print options.
+    print()
+    print('Training options:')
+    print(OmegaConf.to_yaml(args))
+    print()
+    print(f'Output directory:   {args.run_dir}')
+    print(f'Training data:      {args.training_set_kwargs.path}')
+    print(f'Training duration:  {args.total_kimg} kimg')
+    print(f'Number of images:   {args.training_set_kwargs.max_size}')
+    print(f'Image resolution:   {args.training_set_kwargs.resolution}')
+    print(f'Conditional model:  {args.training_set_kwargs.use_labels}')
+    print(f'Dataset x-flips:    {args.training_set_kwargs.xflip}')
+    print()
+
+    # Dry run?
+    if cfg.dry_run:
+        print('Dry run; exiting.')
+        return
+
+    # Create output directory.
+    print('Creating output directory...')
+    if not os.path.exists(args.run_dir):
+        os.makedirs(args.run_dir)
+        with open(os.path.join(args.run_dir, 'training_options.yaml'), 'wt') as fp:
+            OmegaConf.save(config=args, f=fp.name)
+
+    # Launch processes.    
+    print('Launching processes...')
+    if (args.launcher == 'spawn') and (args.num_gpus > 1):
+        args.dist_url = distributed_utils.get_init_file().as_uri()
+        torch.multiprocessing.set_start_method('spawn')
+        torch.multiprocessing.spawn(fn=subprocess_fn, args=(args,), nprocs=args.num_gpus)
+    else:
+        subprocess_fn(rank=0, args=args)
+
+#----------------------------------------------------------------------------
+
+if __name__ == "__main__":
+    if os.getenv('SLURM_ARGS') is not None:
+        # deparcated launcher for slurm jobs.
+        slurm_arg = eval(os.getenv('SLURM_ARGS'))
+        all_args = sys.argv[1:]
+        print(slurm_arg)
+        print(all_args)
+
+        from launcher import launch
+        launch(slurm_arg, all_args)
+    
+    else:
+        main() # pylint: disable=no-value-for-parameter
+
+#----------------------------------------------------------------------------

torch_utils/__init__.py

@@ -0,0 +1,9 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+# empty

torch_utils/custom_ops.py

@@ -0,0 +1,157 @@
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import glob
+import hashlib
+import importlib
+import os
+import re
+import shutil
+import uuid
+
+import torch
+import torch.utils.cpp_extension
+from torch.utils.file_baton import FileBaton
+
+#----------------------------------------------------------------------------
+# Global options.
+
+verbosity = 'brief' # Verbosity level: 'none', 'brief', 'full'
+
+#----------------------------------------------------------------------------
+# Internal helper funcs.
+
+def _find_compiler_bindir():
+    patterns = [
+        'C:/Program Files (x86)/Microsoft Visual Studio/*/Professional/VC/Tools/MSVC/*/bin/Hostx64/x64',
+        'C:/Program Files (x86)/Microsoft Visual Studio/*/BuildTools/VC/Tools/MSVC/*/bin/Hostx64/x64',
+        'C:/Program Files (x86)/Microsoft Visual Studio/*/Community/VC/Tools/MSVC/*/bin/Hostx64/x64',
+        'C:/Program Files (x86)/Microsoft Visual Studio */vc/bin',
+    ]
+    for pattern in patterns:
+        matches = sorted(glob.glob(pattern))
+        if len(matches):
+            return matches[-1]
+    return None
+
+#----------------------------------------------------------------------------
+
+def _get_mangled_gpu_name():
+    name = torch.cuda.get_device_name().lower()
+    out = []
+    for c in name:
+        if re.match('[a-z0-9_-]+', c):
+            out.append(c)
+        else:
+            out.append('-')
+    return ''.join(out)
+
+#----------------------------------------------------------------------------
+# Main entry point for compiling and loading C++/CUDA plugins.
+
+_cached_plugins = dict()
+
+def get_plugin(module_name, sources, headers=None, source_dir=None, **build_kwargs):
+    assert verbosity in ['none', 'brief', 'full']
+    if headers is None:
+        headers = []
+    if source_dir is not None:
+        sources = [os.path.join(source_dir, fname) for fname in sources]
+        headers = [os.path.join(source_dir, fname) for fname in headers]
+
+    # Already cached?
+    if module_name in _cached_plugins:
+        return _cached_plugins[module_name]
+
+    # Print status.
+    if verbosity == 'full':
+        print(f'Setting up PyTorch plugin "{module_name}"...')
+    elif verbosity == 'brief':
+        print(f'Setting up PyTorch plugin "{module_name}"... ', end='', flush=True)
+    verbose_build = (verbosity == 'full')
+
+    # Compile and load.
+    try: # pylint: disable=too-many-nested-blocks
+        # Make sure we can find the necessary compiler binaries.
+        if os.name == 'nt' and os.system("where cl.exe >nul 2>nul") != 0:
+            compiler_bindir = _find_compiler_bindir()
+            if compiler_bindir is None:
+                raise RuntimeError(f'Could not find MSVC/GCC/CLANG installation on this computer. Check _find_compiler_bindir() in "{__file__}".')
+            os.environ['PATH'] += ';' + compiler_bindir
+
+        # Some containers set TORCH_CUDA_ARCH_LIST to a list that can either
+        # break the build or unnecessarily restrict what's available to nvcc.
+        # Unset it to let nvcc decide based on what's available on the
+        # machine.
+        os.environ['TORCH_CUDA_ARCH_LIST'] = ''
+
+        # Incremental build md5sum trickery.  Copies all the input source files
+        # into a cached build directory under a combined md5 digest of the input
+        # source files.  Copying is done only if the combined digest has changed.
+        # This keeps input file timestamps and filenames the same as in previous
+        # extension builds, allowing for fast incremental rebuilds.
+        #
+        # This optimization is done only in case all the source files reside in
+        # a single directory (just for simplicity) and if the TORCH_EXTENSIONS_DIR
+        # environment variable is set (we take this as a signal that the user
+        # actually cares about this.)
+        #
+        # EDIT: We now do it regardless of TORCH_EXTENSIOS_DIR, in order to work
+        # around the *.cu dependency bug in ninja config.
+        #
+        all_source_files = sorted(sources + headers)
+        all_source_dirs = set(os.path.dirname(fname) for fname in all_source_files)
+        if len(all_source_dirs) == 1: # and ('TORCH_EXTENSIONS_DIR' in os.environ):
+
+            # Compute combined hash digest for all source files.
+            hash_md5 = hashlib.md5()
+            for src in all_source_files:
+                with open(src, 'rb') as f:
+                    hash_md5.update(f.read())
+
+            # Select cached build directory name.
+            source_digest = hash_md5.hexdigest()
+            build_top_dir = torch.utils.cpp_extension._get_build_directory(module_name, verbose=verbose_build) # pylint: disable=protected-access
+            cached_build_dir = os.path.join(build_top_dir, f'{source_digest}-{_get_mangled_gpu_name()}')
+
+            if not os.path.isdir(cached_build_dir):
+                tmpdir = f'{build_top_dir}/srctmp-{uuid.uuid4().hex}'
+                os.makedirs(tmpdir)
+                for src in all_source_files:
+                    shutil.copyfile(src, os.path.join(tmpdir, os.path.basename(src)))
+                try:
+                    os.replace(tmpdir, cached_build_dir) # atomic
+                except OSError:
+                    # source directory already exists, delete tmpdir and its contents.
+                    shutil.rmtree(tmpdir)
+                    if not os.path.isdir(cached_build_dir): raise
+
+            # Compile.
+            cached_sources = [os.path.join(cached_build_dir, os.path.basename(fname)) for fname in sources]
+            torch.utils.cpp_extension.load(name=module_name, build_directory=cached_build_dir,
+                verbose=verbose_build, sources=cached_sources, **build_kwargs)
+        else:
+            torch.utils.cpp_extension.load(name=module_name, verbose=verbose_build, sources=sources, **build_kwargs)
+
+        # Load.
+        module = importlib.import_module(module_name)
+
+    except:
+        if verbosity == 'brief':
+            print('Failed!')
+        raise
+
+    # Print status and add to cache dict.
+    if verbosity == 'full':
+        print(f'Done setting up PyTorch plugin "{module_name}".')
+    elif verbosity == 'brief':
+        print('Done.')
+    _cached_plugins[module_name] = module
+    return module
+
+#----------------------------------------------------------------------------

torch_utils/distributed_utils.py

@@ -0,0 +1,213 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+import logging
+import os
+import pickle
+import random
+import socket
+import struct
+import subprocess
+import warnings
+import tempfile
+import uuid
+
+
+from datetime import date
+from pathlib import Path
+from collections import OrderedDict
+from typing import Any, Dict, Mapping
+
+import torch
+import torch.distributed as dist
+
+
+logger = logging.getLogger(__name__)
+
+
+def is_master(args):
+    return args.distributed_rank == 0
+
+
+def init_distributed_mode(rank, args):
+    if "WORLD_SIZE" in os.environ:
+        args.world_size = int(os.environ["WORLD_SIZE"])
+    
+    if args.launcher == 'spawn':  # single node with multiprocessing.spawn
+        args.world_size = args.num_gpus
+        args.rank = rank
+        args.gpu = rank
+    
+    elif 'RANK' in os.environ:
+        args.rank = int(os.environ["RANK"])
+        args.gpu = int(os.environ['LOCAL_RANK'])
+    
+    elif 'SLURM_PROCID' in os.environ:
+        args.rank = int(os.environ['SLURM_PROCID'])
+        args.gpu = args.rank % torch.cuda.device_count()
+    
+    if args.world_size == 1:
+        return
+
+    if 'MASTER_ADDR' in os.environ:
+        args.dist_url = 'tcp://{}:{}'.format(os.environ['MASTER_ADDR'], os.environ['MASTER_PORT'])
+
+    print(f'gpu={args.gpu}, rank={args.rank}, world_size={args.world_size}')
+    args.distributed = True
+    torch.cuda.set_device(args.gpu)
+    args.dist_backend = 'nccl'
+    print('| distributed init (rank {}): {}'.format(args.rank, args.dist_url), flush=True)
+    
+    torch.distributed.init_process_group(backend=args.dist_backend, init_method=args.dist_url,
+                                         world_size=args.world_size, rank=args.rank)
+    torch.distributed.barrier()
+
+
+def gather_list_and_concat(tensor):
+    gather_t = [torch.ones_like(tensor) for _ in range(dist.get_world_size())]
+    dist.all_gather(gather_t, tensor)
+    return torch.cat(gather_t)
+
+
+def get_rank():
+    return dist.get_rank()
+
+
+def get_world_size():
+    return dist.get_world_size()
+
+
+def get_default_group():
+    return dist.group.WORLD
+
+
+def all_gather_list(data, group=None, max_size=16384):
+    """Gathers arbitrary data from all nodes into a list.
+
+    Similar to :func:`~torch.distributed.all_gather` but for arbitrary Python
+    data. Note that *data* must be picklable.
+
+    Args:
+        data (Any): data from the local worker to be gathered on other workers
+        group (optional): group of the collective
+        max_size (int, optional): maximum size of the data to be gathered
+            across workers
+    """
+    rank = get_rank()
+    world_size = get_world_size()
+
+    buffer_size = max_size * world_size
+    if not hasattr(all_gather_list, '_buffer') or \
+            all_gather_list._buffer.numel() < buffer_size:
+        all_gather_list._buffer = torch.cuda.ByteTensor(buffer_size)
+        all_gather_list._cpu_buffer = torch.ByteTensor(max_size).pin_memory()
+    buffer = all_gather_list._buffer
+    buffer.zero_()
+    cpu_buffer = all_gather_list._cpu_buffer
+
+    data = data.cpu()
+    enc = pickle.dumps(data)
+    enc_size = len(enc)
+    header_size = 4  # size of header that contains the length of the encoded data
+    size = header_size + enc_size
+    if size > max_size:
+        raise ValueError('encoded data size ({}) exceeds max_size ({})'.format(size, max_size))
+
+    header = struct.pack(">I", enc_size)
+    cpu_buffer[:size] = torch.ByteTensor(list(header + enc))
+    start = rank * max_size
+    buffer[start:start + size].copy_(cpu_buffer[:size])
+
+    all_reduce(buffer, group=group)
+
+    buffer = buffer.cpu()
+    try:
+        result = []
+        for i in range(world_size):
+            out_buffer = buffer[i * max_size:(i + 1) * max_size]
+            enc_size, = struct.unpack(">I", bytes(out_buffer[:header_size].tolist()))
+            if enc_size > 0:
+                result.append(pickle.loads(bytes(out_buffer[header_size:header_size + enc_size].tolist())))
+        return result
+    except pickle.UnpicklingError:
+        raise Exception(
+            'Unable to unpickle data from other workers. all_gather_list requires all '
+            'workers to enter the function together, so this error usually indicates '
+            'that the workers have fallen out of sync somehow. Workers can fall out of '
+            'sync if one of them runs out of memory, or if there are other conditions '
+            'in your training script that can cause one worker to finish an epoch '
+            'while other workers are still iterating over their portions of the data. '
+            'Try rerunning with --ddp-backend=no_c10d and see if that helps.'
+        )
+
+
+def all_reduce_dict(
+    data: Mapping[str, Any],
+    device,
+    group=None,
+) -> Dict[str, Any]:
+    """
+    AllReduce a dictionary of values across workers. We separately
+    reduce items that are already on the device and items on CPU for
+    better performance.
+
+    Args:
+        data (Mapping[str, Any]): dictionary of data to all-reduce, but
+            cannot be a nested dictionary
+        device (torch.device): device for the reduction
+        group (optional): group of the collective
+    """
+    data_keys = list(data.keys())
+
+    # We want to separately reduce items that are already on the
+    # device and items on CPU for performance reasons.
+    cpu_data = OrderedDict()
+    device_data = OrderedDict()
+    for k in data_keys:
+        t = data[k]
+        if not torch.is_tensor(t):
+            cpu_data[k] = torch.tensor(t, dtype=torch.double)
+        elif t.device.type != device.type:
+            cpu_data[k] = t.to(dtype=torch.double)
+        else:
+            device_data[k] = t.to(dtype=torch.double)
+
+    def _all_reduce_dict(data: OrderedDict):
+        if len(data) == 0:
+            return data
+        buf = torch.stack(list(data.values())).to(device=device)
+        all_reduce(buf, group=group)
+        return {k: buf[i] for i, k in enumerate(data)}
+
+    cpu_data = _all_reduce_dict(cpu_data)
+    device_data = _all_reduce_dict(device_data)
+
+    def get_from_stack(key):
+        if key in cpu_data:
+            return cpu_data[key]
+        elif key in device_data:
+            return device_data[key]
+        raise KeyError
+
+    return OrderedDict([(key, get_from_stack(key)) for key in data_keys])
+
+
+def get_shared_folder() -> Path:
+    user = os.getenv("USER")
+    if Path("/checkpoint/").is_dir():
+        p = Path(f"/checkpoint/{user}/experiments")
+        p.mkdir(exist_ok=True)
+        return p
+    else:
+        p = Path(f"/tmp/experiments")
+        p.mkdir(exist_ok=True)
+        return p
+
+
+def get_init_file():
+    # Init file must not exist, but it's parent dir must exist.
+    os.makedirs(str(get_shared_folder()), exist_ok=True)
+    init_file = Path(str(get_shared_folder()) + f"/{uuid.uuid4().hex}_init")
+    if init_file.exists():
+        os.remove(str(init_file))
+    return init_file
+

torch_utils/misc.py

@@ -0,0 +1,303 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+import re
+import contextlib
+import numpy as np
+import torch
+import warnings
+import dnnlib
+
+#----------------------------------------------------------------------------
+# Cached construction of constant tensors. Avoids CPU=>GPU copy when the
+# same constant is used multiple times.
+
+_constant_cache = dict()
+
+def constant(value, shape=None, dtype=None, device=None, memory_format=None):
+    value = np.asarray(value)
+    if shape is not None:
+        shape = tuple(shape)
+    if dtype is None:
+        dtype = torch.get_default_dtype()
+    if device is None:
+        device = torch.device('cpu')
+    if memory_format is None:
+        memory_format = torch.contiguous_format
+
+    key = (value.shape, value.dtype, value.tobytes(), shape, dtype, device, memory_format)
+    tensor = _constant_cache.get(key, None)
+    if tensor is None:
+        tensor = torch.as_tensor(value.copy(), dtype=dtype, device=device)
+        if shape is not None:
+            tensor, _ = torch.broadcast_tensors(tensor, torch.empty(shape))
+        tensor = tensor.contiguous(memory_format=memory_format)
+        _constant_cache[key] = tensor
+    return tensor
+
+#----------------------------------------------------------------------------
+# Replace NaN/Inf with specified numerical values.
+
+try:
+    nan_to_num = torch.nan_to_num # 1.8.0a0
+except AttributeError:
+    def nan_to_num(input, nan=0.0, posinf=None, neginf=None, *, out=None): # pylint: disable=redefined-builtin
+        assert isinstance(input, torch.Tensor)
+        if posinf is None:
+            posinf = torch.finfo(input.dtype).max
+        if neginf is None:
+            neginf = torch.finfo(input.dtype).min
+        assert nan == 0
+        return torch.clamp(input.unsqueeze(0).nansum(0), min=neginf, max=posinf, out=out)
+
+#----------------------------------------------------------------------------
+# Symbolic assert.
+
+try:
+    symbolic_assert = torch._assert # 1.8.0a0 # pylint: disable=protected-access
+except AttributeError:
+    symbolic_assert = torch.Assert # 1.7.0
+
+#----------------------------------------------------------------------------
+# Context manager to temporarily suppress known warnings in torch.jit.trace().
+# Note: Cannot use catch_warnings because of https://bugs.python.org/issue29672
+
+@contextlib.contextmanager
+def suppress_tracer_warnings():
+    flt = ('ignore', None, torch.jit.TracerWarning, None, 0)
+    warnings.filters.insert(0, flt)
+    yield
+    warnings.filters.remove(flt)
+
+#----------------------------------------------------------------------------
+# Assert that the shape of a tensor matches the given list of integers.
+# None indicates that the size of a dimension is allowed to vary.
+# Performs symbolic assertion when used in torch.jit.trace().
+
+def assert_shape(tensor, ref_shape):
+    if tensor.ndim != len(ref_shape):
+        raise AssertionError(f'Wrong number of dimensions: got {tensor.ndim}, expected {len(ref_shape)}')
+    for idx, (size, ref_size) in enumerate(zip(tensor.shape, ref_shape)):
+        if ref_size is None:
+            pass
+        elif isinstance(ref_size, torch.Tensor):
+            with suppress_tracer_warnings(): # as_tensor results are registered as constants
+                symbolic_assert(torch.equal(torch.as_tensor(size), ref_size), f'Wrong size for dimension {idx}')
+        elif isinstance(size, torch.Tensor):
+            with suppress_tracer_warnings(): # as_tensor results are registered as constants
+                symbolic_assert(torch.equal(size, torch.as_tensor(ref_size)), f'Wrong size for dimension {idx}: expected {ref_size}')
+        elif size != ref_size:
+            raise AssertionError(f'Wrong size for dimension {idx}: got {size}, expected {ref_size}')
+
+#----------------------------------------------------------------------------
+# Function decorator that calls torch.autograd.profiler.record_function().
+
+def profiled_function(fn):
+    def decorator(*args, **kwargs):
+        with torch.autograd.profiler.record_function(fn.__name__):
+            return fn(*args, **kwargs)
+    decorator.__name__ = fn.__name__
+    return decorator
+
+#----------------------------------------------------------------------------
+# Sampler for torch.utils.data.DataLoader that loops over the dataset
+# indefinitely, shuffling items as it goes.
+
+class InfiniteSampler(torch.utils.data.Sampler):
+    def __init__(self, dataset, rank=0, num_replicas=1, shuffle=True, seed=0, window_size=0.5):
+        assert len(dataset) > 0
+        assert num_replicas > 0
+        assert 0 <= rank < num_replicas
+        assert 0 <= window_size <= 1
+        super().__init__(dataset)
+        self.dataset = dataset
+        self.rank = rank
+        self.num_replicas = num_replicas
+        self.shuffle = shuffle
+        self.seed = seed
+        self.window_size = window_size
+
+    def __iter__(self):
+        order = np.arange(len(self.dataset))
+        rnd = None
+        window = 0
+        if self.shuffle:
+            rnd = np.random.RandomState(self.seed)
+            rnd.shuffle(order)
+            window = int(np.rint(order.size * self.window_size))
+
+        idx = 0
+        while True:
+            i = idx % order.size
+            if idx % self.num_replicas == self.rank:
+                yield order[i]
+            if window >= 2:
+                j = (i - rnd.randint(window)) % order.size
+                order[i], order[j] = order[j], order[i]
+            idx += 1
+
+#----------------------------------------------------------------------------
+# Utilities for operating with torch.nn.Module parameters and buffers.
+
+def params_and_buffers(module):
+    assert isinstance(module, torch.nn.Module)
+    return list(module.parameters()) + list(module.buffers())
+
+def named_params_and_buffers(module):
+    assert isinstance(module, torch.nn.Module)
+    return list(module.named_parameters()) + list(module.named_buffers())
+
+def copy_params_and_buffers(src_module, dst_module, require_all=False):
+    assert isinstance(src_module, torch.nn.Module)
+    assert isinstance(dst_module, torch.nn.Module)
+    src_tensors = dict(named_params_and_buffers(src_module))
+    for name, tensor in named_params_and_buffers(dst_module):
+        assert (name in src_tensors) or (not require_all)
+        if name in src_tensors:
+            try:
+                tensor.copy_(src_tensors[name].detach()).requires_grad_(tensor.requires_grad)
+            except Exception as e:
+                print(f'Error loading: {name} {src_tensors[name].shape} {tensor.shape}')
+                raise e
+#----------------------------------------------------------------------------
+# Context manager for easily enabling/disabling DistributedDataParallel
+# synchronization.
+
+@contextlib.contextmanager
+def ddp_sync(module, sync):
+    assert isinstance(module, torch.nn.Module)
+    if sync or not isinstance(module, torch.nn.parallel.DistributedDataParallel):
+        yield
+    else:
+        with module.no_sync():
+            yield
+
+#----------------------------------------------------------------------------
+# Check DistributedDataParallel consistency across processes.
+
+def check_ddp_consistency(module, ignore_regex=None):
+    assert isinstance(module, torch.nn.Module)
+    for name, tensor in named_params_and_buffers(module):
+        fullname = type(module).__name__ + '.' + name
+        if ignore_regex is not None and re.fullmatch(ignore_regex, fullname):
+            continue
+        tensor = tensor.detach()
+        if tensor.is_floating_point():
+            tensor = nan_to_num(tensor)
+        other = tensor.clone()
+        torch.distributed.broadcast(tensor=other, src=0)
+        assert (tensor == other).all(), fullname
+
+#----------------------------------------------------------------------------
+# Print summary table of module hierarchy.
+
+def print_module_summary(module, inputs, max_nesting=3, skip_redundant=True):
+    assert isinstance(module, torch.nn.Module)
+    assert not isinstance(module, torch.jit.ScriptModule)
+    assert isinstance(inputs, (tuple, list))
+
+    # Register hooks.
+    entries = []
+    nesting = [0]
+    def pre_hook(_mod, _inputs):
+        nesting[0] += 1
+    def post_hook(mod, _inputs, outputs):
+        nesting[0] -= 1
+        if nesting[0] <= max_nesting:
+            outputs = list(outputs) if isinstance(outputs, (tuple, list)) else [outputs]
+            outputs = [t for t in outputs if isinstance(t, torch.Tensor)]
+            entries.append(dnnlib.EasyDict(mod=mod, outputs=outputs))
+    hooks = [mod.register_forward_pre_hook(pre_hook) for mod in module.modules()]
+    hooks += [mod.register_forward_hook(post_hook) for mod in module.modules()]
+
+    # Run module.
+    outputs = module(*inputs)
+    for hook in hooks:
+        hook.remove()
+
+    # Identify unique outputs, parameters, and buffers.
+    tensors_seen = set()
+    for e in entries:
+        e.unique_params = [t for t in e.mod.parameters() if id(t) not in tensors_seen]
+        e.unique_buffers = [t for t in e.mod.buffers() if id(t) not in tensors_seen]
+        e.unique_outputs = [t for t in e.outputs if id(t) not in tensors_seen]
+        tensors_seen |= {id(t) for t in e.unique_params + e.unique_buffers + e.unique_outputs}
+
+    # Filter out redundant entries.
+    if skip_redundant:
+        entries = [e for e in entries if len(e.unique_params) or len(e.unique_buffers) or len(e.unique_outputs)]
+
+    # Construct table.
+    rows = [[type(module).__name__, 'Parameters', 'Buffers', 'Output shape', 'Datatype']]
+    rows += [['---'] * len(rows[0])]
+    param_total = 0
+    buffer_total = 0
+    submodule_names = {mod: name for name, mod in module.named_modules()}
+    for e in entries:
+        name = '<top-level>' if e.mod is module else submodule_names[e.mod]
+        param_size = sum(t.numel() for t in e.unique_params)
+        buffer_size = sum(t.numel() for t in e.unique_buffers)
+        output_shapes = [str(list(t.shape)) for t in e.outputs]
+        output_dtypes = [str(t.dtype).split('.')[-1] for t in e.outputs]
+        rows += [[
+            name + (':0' if len(e.outputs) >= 2 else ''),
+            str(param_size) if param_size else '-',
+            str(buffer_size) if buffer_size else '-',
+            (output_shapes + ['-'])[0],
+            (output_dtypes + ['-'])[0],
+        ]]
+        for idx in range(1, len(e.outputs)):
+            rows += [[name + f':{idx}', '-', '-', output_shapes[idx], output_dtypes[idx]]]
+        param_total += param_size
+        buffer_total += buffer_size
+    rows += [['---'] * len(rows[0])]
+    rows += [['Total', str(param_total), str(buffer_total), '-', '-']]
+
+    # Print table.
+    widths = [max(len(cell) for cell in column) for column in zip(*rows)]
+    print()
+    for row in rows:
+        print('  '.join(cell + ' ' * (width - len(cell)) for cell, width in zip(row, widths)))
+    print()
+    return outputs
+
+#----------------------------------------------------------------------------
+
+def get_ddp_func(m, func_name):
+    if hasattr(m, func_name):
+        return getattr(m, func_name)
+    if hasattr(m.module, func_name):
+        return getattr(m.module, func_name)
+    return None
+
+
+#----------------------------------------------------------------------------
+
+@contextlib.contextmanager
+def cuda_time(prefix=""):
+    start = torch.cuda.Event(enable_timing=True)
+    end   = torch.cuda.Event(enable_timing=True)
+    start.record()
+    try:
+        yield 
+    finally: 
+        end.record()
+        torch.cuda.synchronize()
+        print(f'{prefix}: {start.elapsed_time(end)} ms')
+        
+# ---------------------------------------------------------------------------
+
+def get_func(m, f):
+    if hasattr(m, f):
+        return getattr(m, f)
+    elif hasattr(m.module, f):
+        return getattr(m.module, f)
+    else:
+        raise NotImplementedError