Upload 194 files

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+stylegan3-fun/docs/stylegan3-teaser-1920x1006.png filter=lfs diff=lfs merge=lfs -text
+stylegan3-fun/docs/visualizer_screen0.png filter=lfs diff=lfs merge=lfs -text
+stylegan3-fun/out/seed0002.png filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,87 @@
+import os
+import sys
+import re
+from typing import List, Optional, Tuple, Union
+import random
+
+sys.path.append('stylegan3-fun')  # change this to the path where dnnlib is located
+
+import numpy as np
+import PIL.Image
+import torch
+import streamlit as st
+import dnnlib
+import legacy
+
+
+def parse_range(s: Union[str, List]) -> List[int]:
+    '''Parse a comma separated list of numbers or ranges and return a list of ints.
+
+    Example: '1,2,5-10' returns [1, 2, 5, 6, 7]
+    '''
+    if isinstance(s, list): return s
+    ranges = []
+    range_re = re.compile(r'^(\d+)-(\d+)$')
+    for p in s.split(','):
+        m = range_re.match(p)
+        if m:
+            ranges.extend(range(int(m.group(1)), int(m.group(2))+1))
+        else:
+            ranges.append(int(p))
+    return ranges
+
+def make_transform(translate: Tuple[float,float], angle: float):
+    m = np.eye(3)
+    s = np.sin(angle/360.0*np.pi*2)
+    c = np.cos(angle/360.0*np.pi*2)
+    m[0][0] = c
+    m[0][1] = s
+    m[0][2] = translate[0]
+    m[1][0] = -s
+    m[1][1] = c
+    m[1][2] = translate[1]
+    return m
+
+def generate_image(network_pkl: str, seed: int, truncation_psi: float, noise_mode: str, translate: Tuple[float,float], rotate: float, class_idx: Optional[int]):
+    print('Loading networks from "%s"...' % network_pkl)
+    device = torch.device('cuda')
+    with open(network_pkl, 'rb') as f:
+        G = legacy.load_network_pkl(f)['G_ema'].to(device) # type: ignore
+
+    # Labels.
+    label = torch.zeros([1, G.c_dim], device=device)
+    if G.c_dim != 0:
+        if class_idx is None:
+            raise Exception('Must specify class label when using a conditional network')
+        label[:, class_idx] = 1
+
+    z = torch.from_numpy(np.random.RandomState(seed).randn(1, G.z_dim)).to(device)
+
+    if hasattr(G.synthesis, 'input'):
+        m = make_transform(translate, rotate)
+        m = np.linalg.inv(m)
+        G.synthesis.input.transform.copy_(torch.from_numpy(m))
+
+    img = G(z, label, truncation_psi=truncation_psi, noise_mode=noise_mode)
+    img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)
+    img = PIL.Image.fromarray(img[0].cpu().numpy(), 'RGB')
+    return img
+
+def main():
+    st.title('Kpop Face Generator')
+
+    st.write('Press the button below to generate a new image:')
+    if st.button('Generate'):
+        network_pkl = 'kpopGG.pkl'
+        seed = random.randint(0, 99999)
+        truncation_psi = 0.45
+        noise_mode = 'const'
+        translate = (0.0, 0.0)
+        rotate = 0.0
+        class_idx = None
+
+        image = generate_image(network_pkl, seed, truncation_psi, noise_mode, translate, rotate, class_idx)
+        st.image(image)
+
+if __name__ == "__main__":
+    main()

stylegan3-fun/.github/FUNDING.yml

@@ -0,0 +1,3 @@
+# These are supported funding model platforms
+
+github: PDillis # Replace with up to 4 GitHub Sponsors-enabled usernames e.g., [user1, user2]

stylegan3-fun/.github/ISSUE_TEMPLATE/bug_report.md

@@ -0,0 +1,35 @@
+---
+name: Bug report
+about: Create a report to help us improve
+title: ''
+labels: ''
+assignees: ''
+
+---
+
+**Describe the bug**
+A clear and concise description of what the bug is.
+
+**To Reproduce**
+Steps to reproduce the behavior:
+1. In '...' directory, run command '...'
+2. See error (copy&paste full log, including exceptions and **stacktraces**).
+
+Please copy&paste text instead of screenshots for better searchability.
+
+**Expected behavior**
+A clear and concise description of what you expected to happen.
+
+**Screenshots**
+If applicable, add screenshots to help explain your problem.
+
+**Desktop (please complete the following information):**
+ - OS: [e.g. Linux Ubuntu 20.04, Windows 10]
+ - PyTorch version (e.g., pytorch 1.9.0)
+ - CUDA toolkit version (e.g., CUDA 11.4)
+ - NVIDIA driver version
+ - GPU [e.g., Titan V, RTX 3090]
+ - Docker: did you use Docker?  If yes, specify docker image URL (e.g., nvcr.io/nvidia/pytorch:21.08-py3)
+
+**Additional context**
+Add any other context about the problem here.

stylegan3-fun/.github/ISSUE_TEMPLATE/feature_request.md

@@ -0,0 +1,20 @@
+---
+name: Feature request
+about: Suggest an idea for this project
+title: ''
+labels: ''
+assignees: ''
+
+---
+
+**Is your feature request related to a problem? Please describe.**
+A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]
+
+**Describe the solution you'd like**
+A clear and concise description of what you want to happen.
+
+**Describe alternatives you've considered**
+A clear and concise description of any alternative solutions or features you've considered.
+
+**Additional context**
+Add any other context or screenshots about the feature request here.

stylegan3-fun/.gitignore

@@ -0,0 +1,173 @@
+# Byte-compiled / optimized / DLL files
+*/**/__pycache__/
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+.idea/
+
+# Conda temp
+.condatmp/
+
+# SGAN specific folders
+datasets/
+dlatents/
+out/
+training-runs/
+pretrained/
+video/
+_screenshots/

stylegan3-fun/Dockerfile

@@ -0,0 +1,19 @@
+# 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.
+
+FROM nvcr.io/nvidia/pytorch:21.08-py3
+
+ENV PYTHONDONTWRITEBYTECODE 1
+ENV PYTHONUNBUFFERED 1
+
+RUN pip install imageio imageio-ffmpeg==0.4.4 pyspng==0.1.0
+
+WORKDIR /workspace
+
+RUN (printf '#!/bin/bash\nexec \"$@\"\n' >> /entry.sh) && chmod a+x /entry.sh
+ENTRYPOINT ["/entry.sh"]

stylegan3-fun/LICENSE.txt

@@ -0,0 +1,97 @@
+Copyright (c) 2021, NVIDIA Corporation & affiliates. All rights reserved.
+
+
+NVIDIA Source Code License for StyleGAN3
+
+
+=======================================================================
+
+1. Definitions
+
+"Licensor" means any person or entity that distributes its Work.
+
+"Software" means the original work of authorship made available under
+this License.
+
+"Work" means the Software and any additions to or derivative works of
+the Software that are made available under this License.
+
+The terms "reproduce," "reproduction," "derivative works," and
+"distribution" have the meaning as provided under U.S. copyright law;
+provided, however, that for the purposes of this License, derivative
+works shall not include works that remain separable from, or merely
+link (or bind by name) to the interfaces of, the Work.
+
+Works, including the Software, are "made available" under this License
+by including in or with the Work either (a) a copyright notice
+referencing the applicability of this License to the Work, or (b) a
+copy of this License.
+
+2. License Grants
+
+    2.1 Copyright Grant. Subject to the terms and conditions of this
+    License, each Licensor grants to you a perpetual, worldwide,
+    non-exclusive, royalty-free, copyright license to reproduce,
+    prepare derivative works of, publicly display, publicly perform,
+    sublicense and distribute its Work and any resulting derivative
+    works in any form.
+
+3. Limitations
+
+    3.1 Redistribution. You may reproduce or distribute the Work only
+    if (a) you do so under this License, (b) you include a complete
+    copy of this License with your distribution, and (c) you retain
+    without modification any copyright, patent, trademark, or
+    attribution notices that are present in the Work.
+
+    3.2 Derivative Works. You may specify that additional or different
+    terms apply to the use, reproduction, and distribution of your
+    derivative works of the Work ("Your Terms") only if (a) Your Terms
+    provide that the use limitation in Section 3.3 applies to your
+    derivative works, and (b) you identify the specific derivative
+    works that are subject to Your Terms. Notwithstanding Your Terms,
+    this License (including the redistribution requirements in Section
+    3.1) will continue to apply to the Work itself.
+
+    3.3 Use Limitation. The Work and any derivative works thereof only
+    may be used or intended for use non-commercially. Notwithstanding
+    the foregoing, NVIDIA and its affiliates may use the Work and any
+    derivative works commercially. As used herein, "non-commercially"
+    means for research or evaluation purposes only.
+
+    3.4 Patent Claims. If you bring or threaten to bring a patent claim
+    against any Licensor (including any claim, cross-claim or
+    counterclaim in a lawsuit) to enforce any patents that you allege
+    are infringed by any Work, then your rights under this License from
+    such Licensor (including the grant in Section 2.1) will terminate
+    immediately.
+
+    3.5 Trademarks. This License does not grant any rights to use any
+    Licensor’s or its affiliates’ names, logos, or trademarks, except
+    as necessary to reproduce the notices described in this License.
+
+    3.6 Termination. If you violate any term of this License, then your
+    rights under this License (including the grant in Section 2.1) will
+    terminate immediately.
+
+4. Disclaimer of Warranty.
+
+THE WORK IS PROVIDED "AS IS" WITHOUT WARRANTIES OR CONDITIONS OF ANY
+KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WARRANTIES OR CONDITIONS OF
+MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, TITLE OR
+NON-INFRINGEMENT. YOU BEAR THE RISK OF UNDERTAKING ANY ACTIVITIES UNDER
+THIS LICENSE.
+
+5. Limitation of Liability.
+
+EXCEPT AS PROHIBITED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL
+THEORY, WHETHER IN TORT (INCLUDING NEGLIGENCE), CONTRACT, OR OTHERWISE
+SHALL ANY LICENSOR BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY DIRECT,
+INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES ARISING OUT OF
+OR RELATED TO THIS LICENSE, THE USE OR INABILITY TO USE THE WORK
+(INCLUDING BUT NOT LIMITED TO LOSS OF GOODWILL, BUSINESS INTERRUPTION,
+LOST PROFITS OR DATA, COMPUTER FAILURE OR MALFUNCTION, OR ANY OTHER
+COMMERCIAL DAMAGES OR LOSSES), EVEN IF THE LICENSOR HAS BEEN ADVISED OF
+THE POSSIBILITY OF SUCH DAMAGES.
+
+=======================================================================

stylegan3-fun/README.md

@@ -0,0 +1,507 @@
+# StyleGAN3-Fun<br><sub>Let's have fun with StyleGAN2/ADA/3!</sub>
+
+SOTA GANs are hard to train and to explore, and StyleGAN2/ADA/3 are no different. The point of this repository is to allow
+the user to both easily train and explore the trained models without unnecessary headaches. 
+
+As [before](https://github.com/PDillis/stylegan2-fun), we will build upon the official repository, which has the advantage
+of being backwards-compatible. As such, we can use our previously-trained models from StyleGAN2 and StyleGAN2-ADA. Please
+get acquainted with the official repository and its codebase, as we will be building upon it and as such, increase its
+capabilities (but hopefully not its complexity!). 
+
+## Additions
+
+This repository adds/has the following changes (not yet the complete list):
+
+* ***Dataset Setup (`dataset_tool.py`)***
+  * **RGBA support**, so revert saving images to `.png` ([Issue #156](https://github.com/NVlabs/stylegan3/issues/156) by @1378dm). Training can use RGBA and images can be generated.
+    * ***TODO:*** ~~Check that training code is correct for normalizing the alpha channel~~, as well as making the 
+      interpolation code work with this new format (look into [`moviepy.editor.VideoClip`](https://zulko.github.io/moviepy/getting_started/videoclips.html?highlight=mask#mask-clips))
+    * For now, interpolation videos will only be saved in RGB format, e.g., discarding the alpha channel.
+  * **`--center-crop-tall`**: add vertical black bars to the sides of each image in the dataset (rectangular images, with height > width),
+    and you wish to train a square model, in the same vein as the horizontal bars added when using `--center-crop-wide` (where width > height).
+    * This is useful when you don't want to lose information from the left and right side of the image by only using the center
+      crop (ibidem for `--center-crop-wide`, but for the top and bottom of the image)
+    * Note that each image doesn't have to be of the same size, and the added bars will only ensure you get a square image, which will then be
+      resized to the model's desired resolution (set by `--resolution`).
+  * Grayscale images in the dataset are converted to `RGB`
+    * If you want to turn this off, remove the respective line in `dataset_tool.py`, e.g., if your dataset is made of images in a folder, then the function to be used is
+    `open_image_folder` in `dataset_tool.py`, and the line to be removed is `img = img.convert('RGB')` in the `iterate_images` inner function.
+  * The dataset can be forced to be of a specific number of channels, that is, grayscale, RGB or RGBA.
+    * To use this, set `--force-channels=1` for grayscale, `--force-channels=3` for RGB, and `--force-channels=4` for RGBA.
+  * If the dataset tool encounters an error, print it along the offending image, but continue with the rest of the dataset 
+    ([PR #39](https://github.com/NVlabs/stylegan3/pull/39) from [Andreas Jansson](https://github.com/andreasjansson)).
+  * For conditional models, we can use the subdirectories as the classes by adding `--subfolders-as-labels`. This will 
+  generate the `dataset.json` file automatically as done by @pbaylies [here](https://github.com/pbaylies/stylegan2-ada/blob/a8f0b1c891312631f870c94f996bcd65e0b8aeef/dataset_tool.py#L772)
+    * Additionally, in the `--source` folder, we will save a `class_labels.txt` file, to further know which classes correspond to each subdirectory.
+
+* ***Training***
+  * Add `--cfg=stylegan2-ext`, which uses @aydao's extended modifications for handling large and diverse datasets. 
+    * A good explanation is found in Gwern's blog [here](https://gwern.net/face#extended-stylegan2-danbooru2019-aydao)
+    * If you wish to fine-tune from @aydao's Anime model, use `--cfg=stylegan2-ext --resume=anime512` when running `train.py`
+    * Note: ***This is an extremely experimental configuration!*** The `.pkl` files will be ~1.1Gb each and training will slow down
+      significantly. Use at your own risk!
+  * `--blur-percent`: Blur both real and generated images before passing them to the Discriminator. 
+    * The blur (`blur_init_sigma=10.0`) will completely fade after the selected percentage of the training is completed (using a linear ramp).
+    * Another experimental feature, should help with datasets that have a lot of variation, and you wish the model to slowly
+      learn to generate the objects and then its details.
+  * `--mirrory`: Added vertical mirroring for doubling the dataset size (quadrupling if `--mirror` is used; make sure your dataset has either or both 
+    of these symmetries in order for it to make sense to use them)
+  * `--gamma`: If no R1 regularization is provided, the heuristic formula from [StyleGAN](https://github.com/NVlabs/stylegan2) will be used.
+    * Specifically, we will set `gamma=0.0002 * resolution ** 2 / batch_size`
+  * `--aug`: ***TODO:*** add [Deceive-D/APA](https://github.com/EndlessSora/DeceiveD) as an option.
+  * `--augpipe`: Now available to use is [StyleGAN2-ADA's](https://github.com/NVlabs/stylegan2-ada-pytorch) full list of augpipe, i.e., individual augmentations (`blit`, `geom`, `color`, `filter`, `noise`, `cutout`) or their combinations (`bg`, `bgc`, `bgcf`, `bgcfn`, `bgcfnc`).
+  * `--img-snap`: Set when to save snapshot images, so now it's independent of when the model is saved (e.g., save image snapshots more often to know how the model is training without saving the model itself, to save space).
+  * `--snap-res`: The resolution of the snapshots, depending on how many images you wish to see per snapshot. Available resolutions: `1080p`, `4k`, and `8k`.
+  * `--resume-kimg`: Starting number of `kimg`, useful when continuing training a previous run
+  * `--outdir`: Automatically set as `training-runs`, so no need to set beforehand (in general this is true throughout the repository)
+  * `--metrics`: Now set by default to `None`, so there's no need to worry about this one
+  * `--freezeD`: Renamed `--freezed` for better readability
+  * `--freezeM`: Freeze the first layers of the Mapping Network Gm (`G.mapping`)
+  * `--freezeE`: Freeze the embedding layer of the Generator (for class-conditional models)
+  * `--freezeG`: ***TODO:*** Freeze the first layers of the Synthesis Network (`G.synthesis`; less cost to transfer learn, focus on high layers?)
+  * `--resume`: All available pre-trained models from NVIDIA (and more) can be used with a simple dictionary, depending on the `--cfg` used.
+  For example, if you wish to use StyleGAN3's `config-r`, then set `--cfg=stylegan3-r`. In addition, if you wish to transfer learn from FFHQU at 1024 resolution, set `--resume=ffhqu1024`.
+    * The full list of currently available models to transfer learn from (or synthesize new images with) is the following (***TODO:*** add small description of each model, 
+    so the user can better know which to use for their particular use-case; proper citation to original authors as well):
+    
+        <details>
+        <summary>StyleGAN2 models</summary>
+    
+      1. Majority, if not all, are `config-f`: set `--cfg=stylegan2`
+         * `ffhq256`
+         * `ffhqu256`
+         * `ffhq512`
+         * `ffhq1024`
+         * `ffhqu1024`
+         * `celebahq256`
+         * `lsundog256`
+         * `afhqcat512`
+         * `afhqdog512`
+         * `afhqwild512`
+         * `afhq512`
+         * `brecahad512`
+         * `cifar10` (conditional, 10 classes)
+         * `metfaces1024`
+         * `metfacesu1024`
+         * `lsuncar512` (config-f)
+         * `lsuncat256` (config-f)
+         * `lsunchurch256` (config-f)
+         * `lsunhorse256` (config-f)
+         * `minecraft1024` (thanks to @jeffheaton)
+         * `imagenet512` (thanks to @shawwn)
+         * `wikiart1024-C` (conditional, 167 classes; thanks to @pbaylies)
+         * `wikiart1024-U` (thanks to @pbaylies)
+         * `maps1024` (thanks to @tjukanov)
+         * `fursona512` (thanks to @arfafax)
+         * `mlpony512` (thanks to @arfafax)
+         * `lhq1024` (thanks to @justinpinkney)
+         * `afhqcat256` (Deceive-D/APA models)
+         * `anime256` (Deceive-D/APA models)
+         * `cub256` (Deceive-D/APA models)
+         * `sddogs1024` (Self-Distilled StyleGAN models)
+         * `sdelephant512` (Self-Distilled StyleGAN models)
+         * `sdhorses512` (Self-Distilled StyleGAN models)
+         * `sdbicycles256` (Self-Distilled StyleGAN models)
+         * `sdlions512` (Self-Distilled StyleGAN models)
+         * `sdgiraffes512` (Self-Distilled StyleGAN models)
+         * `sdparrots512` (Self-Distilled StyleGAN models)
+      2. Extended StyleGAN2 config from @aydao: set `--cfg=stylegan2-ext`
+         * `anime512` (thanks to @aydao; writeup by @gwern: https://gwern.net/Faces#extended-stylegan2-danbooru2019-aydao)
+        </details>
+    
+        <details>
+        <summary>StyleGAN3 models</summary>
+    
+      1. `config-t`: set `--cfg=stylegan3-t`
+         * `afhq512`
+         * `ffhqu256`
+         * `ffhq1024`
+         * `ffhqu1024`
+         * `metfaces1024`
+         * `metfacesu1024`
+         * `landscapes256` (thanks to @justinpinkney)
+         * `wikiart1024` (thanks to @justinpinkney)
+         * `mechfuture256` (thanks to @edstoica; 29 kimg tick)
+         * `vivflowers256` (thanks to @edstoica; 68 kimg tick)
+         * `alienglass256` (thanks to @edstoica; 38 kimg tick)
+         * `scificity256` (thanks to @edstoica; 210 kimg tick)
+         * `scifiship256` (thanks to @edstoica; 168 kimg tick)
+      2. `config-r`: set `--cfg=stylegan3-r`
+         * `afhq512`
+         * `ffhq1024`
+         * `ffhqu1024`
+         * `ffhqu256`
+         * `metfaces1024`
+         * `metfacesu1024`
+        </details>
+
+    * The main sources of these pretrained models are both the [official NVIDIA repository](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/research/models/stylegan3),
+      as well as other community repositories, such as [Justin Pinkney](https://github.com/justinpinkney) 's [Awesome Pretrained StyleGAN2](https://github.com/justinpinkney/awesome-pretrained-stylegan2)
+      and [Awesome Pretrained StyleGAN3](https://github.com/justinpinkney/awesome-pretrained-stylegan3), [Deceive-D/APA](https://github.com/EndlessSora/DeceiveD),
+      [Self-Distilled StyleGAN/Internet Photos](https://github.com/self-distilled-stylegan/self-distilled-internet-photos), and [edstoica](https://github.com/edstoica) 's 
+      [Wombo Dream](https://www.wombo.art/) [-based models](https://github.com/edstoica/lucid_stylegan3_datasets_models). Others can be found around the net and are properly credited in this repository,
+      so long as they can be easily downloaded with [`dnnlib.util.open_url`](https://github.com/PDillis/stylegan3-fun/blob/4ce9d6f7601641ba1e2906ed97f2739a63fb96e2/dnnlib/util.py#L396).
+
+* ***Interpolation videos***
+    * [Random interpolation](https://youtu.be/DNfocO1IOUE)
+      * [Generate images/interpolations with the internal representations of the model](https://nvlabs-fi-cdn.nvidia.com/_web/stylegan3/videos/video_8_internal_activations.mp4)
+        * Usage: Add `--layer=<layer_name>` to specify which layer to use for interpolation.
+        * If you don't know the names of the layers available for your model, add the flag `--available-layers` and the 
+        layers will be printed to the console, along their names, number of channels, and sizes. 
+        * Use one of `--grayscale` or `--rgb` to specify whether to save the images as grayscale or RGB during the interpolation.
+        * For `--rgb`, three consecutive channels (starting at `--starting-channel=0`) will be used to create the RGB image. For `--grayscale`, only the first channel will be used.
+    * Style-mixing
+    * [Sightseeding](https://twitter.com/PDillis/status/1270341433401249793?s=20&t=yLueNkagqsidZFqZ2jNPAw) (jumpiness has been fixed)
+    * [Circular interpolation](https://youtu.be/4nktYGjSVHg)
+    * [Visual-reactive interpolation](https://youtu.be/KoEAkPnE-zA) (Beta)
+    * Audiovisual-reactive interpolation (TODO)
+    * ***TODO:*** Give support to RGBA models!
+* ***Projection into the latent space***
+    * [Project into $\mathcal{W}+$](https://arxiv.org/abs/1904.03189)
+    * Additional losses to use for better projection (e.g., using VGG16 or [CLIP](https://github.com/openai/CLIP))
+* ***[Discriminator Synthesis](https://arxiv.org/abs/2111.02175)*** (official code)
+    * Generate a static image (`python discriminator_synthesis.py dream --help`) or a [video](https://youtu.be/hEJKWL2VQTE) with a feedback loop (`python discriminator_synthesis.py dream-zoom --help`,
+     `python discriminator_synthesis.py channel-zoom --help`, or `python discriminator_synthesis.py interp --help`)
+    * Start from a random image (`random` for noise or `perlin` for 2D fractal Perlin noise, using 
+      [Mathieu Duchesneau's implementation](https://github.com/duchesneaumathieu/pyperlin)) or from an existing one
+* ***Expansion on GUI/`visualizer.py`***
+    * Added the rest of the affine transformations
+    * Added widget for class-conditional models (***TODO:*** mix classes with continuous values for `cls`!)
+* ***General model and code additions***
+    * [Multi-modal truncation trick](https://arxiv.org/abs/2202.12211): find the different clusters in your model and use the closest one to your dlatent, in order to increase the fidelity
+      * Usage: Run `python multimodal_truncation.py get-centroids --network=<path_to_model>` to use default values; for extra options, run `python multimodal_truncation.py get-centroids --help`
+    * StyleGAN3: anchor the latent space for easier to follow interpolations (thanks to [Rivers Have Wings](https://github.com/crowsonkb) and [nshepperd](https://github.com/nshepperd)).
+    * Use CPU instead of GPU if desired (not recommended, but perfectly fine for generating images, whenever the custom CUDA kernels fail to compile).
+    * Add missing dependencies and channels so that the [`conda`](https://docs.conda.io/en/latest/) environment is correctly setup in Windows 
+      (PR's [#111](https://github.com/NVlabs/stylegan3/pull/111)/[#125](https://github.com/NVlabs/stylegan3/pull/125) and [#80](https://github.com/NVlabs/stylegan3/pull/80) /[#143](https://github.com/NVlabs/stylegan3/pull/143) from the base, respectively)
+    * Use [StyleGAN-NADA](https://github.com/rinongal/StyleGAN-nada) models with any part of the code (Issue [#9](https://github.com/PDillis/stylegan3-fun/issues/9))
+      * The StyleGAN-NADA models must first be converted via [Vadim Epstein](https://github.com/eps696) 's conversion code found [here](https://github.com/eps696/stylegan2ada#tweaking-models).
+    * Add PR [#173](https://github.com/NVlabs/stylegan3/pull/173) for adding the last remaining unknown kwarg for using StyleGAN2 models using TF 1.15.
+* ***TODO*** list (this is a long one with more to come, so any help is appreciated):
+  * Add `--device={cuda, ref}` option to pass to each of the custom operations in order to (theoretically) be able to use AMD GPUs, as explained in
+  @l4rz's post [here](https://github.com/l4rz/practical-aspects-of-stylegan2-training#hardware)
+  * Define a [custom Generator](https://github.com/dvschultz/stylegan2-ada-pytorch/blob/59e05bb115c1c7d0de56be0523754076c2b7ee83/legacy.py#L131) in `legacy.py` to modify the output size
+    * Related: the [multi-latent](https://github.com/dvschultz/stylegan2-ada-pytorch/blob/main/training/stylegan2_multi.py), i.e., the one from [@eps696](https://github.com/eps696/stylegan2)
+  * Add [Top-K training](https://arxiv.org/abs/2002.06224) as done [here](https://github.com/dvschultz/stylegan2-ada/blob/8f4ab24f494483542d31bf10f4fdb0005dc62739/train.py#L272) and [here](https://github.com/dvschultz/stylegan2-ada-pytorch/blob/59e05bb115c1c7d0de56be0523754076c2b7ee83/training/loss.py#L79)
+  * Add panorama/SinGAN/feature interpolation from [StyleGAN of All Trades](https://arxiv.org/abs/2111.01619)
+  * [PTI](https://github.com/danielroich/PTI) for better inversion
+  * [Better sampling](https://arxiv.org/abs/2110.08009)
+  * [Add cross-model interpolation](https://twitter.com/arfafax/status/1297681537337446400?s=20&t=xspnTaLFTvd7y4krg8tkxA)
+  * Blend different models (average checkpoints, copy weights, create initial network), as in @aydao's [StyleGAN2-Surgery](https://github.com/aydao/stylegan2-surgery)
+  * Add multi-crop for the dataset creation, as used in [Earth View](https://github.com/PDillis/earthview#multi-crop---data_augmentpy).
+  * Make it easy to download pretrained models from Drive, otherwise a lot of models can't be used with `dnnlib.util.open_url`
+    (e.g., [StyleGAN-Human](https://github.com/stylegan-human/StyleGAN-Human) models)
+  * Finish documentation for better user experience, add videos/images, code samples, visuals...
+  * Add [Ensembling Off-the-shelf Models for GAN Training](https://arxiv.org/abs/2112.09130) and [Any-resolution Training for High-resolution Image Synthesis](https://chail.github.io/anyres-gan/)
+
+## Notebooks (Coming soon!)
+
+## Sponsors ![GitHub Sponsor](https://img.shields.io/github/sponsors/PDillis?label=Sponsor&logo=GitHub)
+
+This repository has been sponsored by:
+
+[isosceles](https://www.jasonfletcher.info/vjloops/)
+
+Thank you so much! 
+
+If you wish to sponsor me, click here: [![](https://img.shields.io/static/v1?label=Sponsor&message=%E2%9D%A4&logo=GitHub&color=%23fe8e86)](https://github.com/sponsors/PDillis)
+
+
+---
+
+## Alias-Free Generative Adversarial Networks (StyleGAN3)<br><sub>Official PyTorch implementation of the NeurIPS 2021 paper</sub>
+
+![Teaser image](./docs/stylegan3-teaser-1920x1006.png)
+
+**Alias-Free Generative Adversarial Networks**<br>
+Tero Karras, Miika Aittala, Samuli Laine, Erik H&auml;rk&ouml;nen, Janne Hellsten, Jaakko Lehtinen, Timo Aila<br>
+https://nvlabs.github.io/stylegan3<br>
+
+Abstract: *We observe that despite their hierarchical convolutional nature, the synthesis process of typical generative adversarial networks depends on absolute pixel coordinates in an unhealthy manner. This manifests itself as, e.g., detail appearing to be glued to image coordinates instead of the surfaces of depicted objects. We trace the root cause to careless signal processing that causes aliasing in the generator network. Interpreting all signals in the network as continuous, we derive generally applicable, small architectural changes that guarantee that unwanted information cannot leak into the hierarchical synthesis process. The resulting networks match the FID of StyleGAN2 but differ dramatically in their internal representations, and they are fully equivariant to translation and rotation even at subpixel scales. Our results pave the way for generative models better suited for video and animation.*
+
+For business inquiries, please visit our website and submit the form: [NVIDIA Research Licensing](https://www.nvidia.com/en-us/research/inquiries/)
+
+## Release notes
+
+This repository is an updated version of [stylegan2-ada-pytorch](https://github.com/NVlabs/stylegan2-ada-pytorch), with several new features:
+- Alias-free generator architecture and training configurations (`stylegan3-t`, `stylegan3-r`).
+- Tools for interactive visualization (`visualizer.py`), spectral analysis (`avg_spectra.py`), and video generation (`gen_video.py`).
+- Equivariance metrics (`eqt50k_int`, `eqt50k_frac`, `eqr50k`).
+- General improvements: reduced memory usage, slightly faster training, bug fixes.
+
+Compatibility:
+- Compatible with old network pickles created using [stylegan2-ada](https://github.com/NVlabs/stylegan2-ada) and [stylegan2-ada-pytorch](https://github.com/NVlabs/stylegan2-ada-pytorch).  (Note: running old StyleGAN2 models on StyleGAN3 code will produce the same results as running them on stylegan2-ada/stylegan2-ada-pytorch.  To benefit from the StyleGAN3 architecture, you need to retrain.)
+- Supports old StyleGAN2 training configurations, including ADA and transfer learning. See [Training configurations](./docs/configs.md) for details.
+- Improved compatibility with Ampere GPUs and newer versions of PyTorch, CuDNN, etc.
+
+## Synthetic image detection
+
+While new generator approaches enable new media synthesis capabilities, they may also present a new challenge for AI forensics algorithms for detection and attribution of synthetic media. In collaboration with digital forensic researchers participating in DARPA's SemaFor program, we curated a synthetic image dataset that allowed the researchers to test and validate the performance of their image detectors in advance of the public release. Please see [here](https://github.com/NVlabs/stylegan3-detector) for more details.
+
+## Additional material
+
+- [Result videos](https://nvlabs-fi-cdn.nvidia.com/stylegan3/videos/)
+- [Curated example images](https://nvlabs-fi-cdn.nvidia.com/stylegan3/images/)
+- [StyleGAN3 pre-trained models](https://ngc.nvidia.com/catalog/models/nvidia:research:stylegan3) for config T (translation equiv.) and config R (translation and rotation equiv.)
+  > <sub>Access individual networks via `https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/<MODEL>`, where `<MODEL>` is one of:</sub><br>
+  > <sub>`stylegan3-t-ffhq-1024x1024.pkl`, `stylegan3-t-ffhqu-1024x1024.pkl`, `stylegan3-t-ffhqu-256x256.pkl`</sub><br>
+  > <sub>`stylegan3-r-ffhq-1024x1024.pkl`, `stylegan3-r-ffhqu-1024x1024.pkl`, `stylegan3-r-ffhqu-256x256.pkl`</sub><br>
+  > <sub>`stylegan3-t-metfaces-1024x1024.pkl`, `stylegan3-t-metfacesu-1024x1024.pkl`</sub><br>
+  > <sub>`stylegan3-r-metfaces-1024x1024.pkl`, `stylegan3-r-metfacesu-1024x1024.pkl`</sub><br>
+  > <sub>`stylegan3-t-afhqv2-512x512.pkl`</sub><br>
+  > <sub>`stylegan3-r-afhqv2-512x512.pkl`</sub><br>
+- [StyleGAN2 pre-trained models](https://ngc.nvidia.com/catalog/models/nvidia:research:stylegan2) compatible with this codebase
+  > <sub>Access individual networks via `https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan2/versions/1/files/<MODEL>`, where `<MODEL>` is one of:</sub><br>
+  > <sub>`stylegan2-ffhq-1024x1024.pkl`, `stylegan2-ffhq-512x512.pkl`, `stylegan2-ffhq-256x256.pkl`</sub><br>
+  > <sub>`stylegan2-ffhqu-1024x1024.pkl`, `stylegan2-ffhqu-256x256.pkl`</sub><br>
+  > <sub>`stylegan2-metfaces-1024x1024.pkl`, `stylegan2-metfacesu-1024x1024.pkl`</sub><br>
+  > <sub>`stylegan2-afhqv2-512x512.pkl`</sub><br>
+  > <sub>`stylegan2-afhqcat-512x512.pkl`, `stylegan2-afhqdog-512x512.pkl`, `stylegan2-afhqwild-512x512.pkl`</sub><br>
+  > <sub>`stylegan2-brecahad-512x512.pkl`, `stylegan2-cifar10-32x32.pkl`</sub><br>
+  > <sub>`stylegan2-celebahq-256x256.pkl`, `stylegan2-lsundog-256x256.pkl`</sub><br>
+
+## Requirements
+
+* Linux and Windows are supported, but we recommend Linux for performance and compatibility reasons.
+* 1&ndash;8 high-end NVIDIA GPUs with at least 12 GB of memory. We have done all testing and development using Tesla V100 and A100 GPUs.
+* 64-bit Python 3.8 and PyTorch 1.9.0 (or later). See https://pytorch.org for PyTorch install instructions.
+* CUDA toolkit 11.1 or later.  (Why is a separate CUDA toolkit installation required?  See [Troubleshooting](./docs/troubleshooting.md#why-is-cuda-toolkit-installation-necessary)).
+* GCC 7 or later (Linux) or Visual Studio (Windows) compilers.  Recommended GCC version depends on CUDA version, see for example [CUDA 11.4 system requirements](https://docs.nvidia.com/cuda/archive/11.4.1/cuda-installation-guide-linux/index.html#system-requirements).
+* Python libraries: see [environment.yml](./environment.yml) for exact library dependencies.  You can use the following commands with Miniconda3 to create and activate your StyleGAN3 Python environment:
+  - `conda env create -f environment.yml`
+  - `conda activate stylegan3`
+* Docker users:
+  - Ensure you have correctly installed the [NVIDIA container runtime](https://docs.docker.com/config/containers/resource_constraints/#gpu).
+  - Use the [provided Dockerfile](./Dockerfile) to build an image with the required library dependencies.
+
+The code relies heavily on custom PyTorch extensions that are compiled on the fly using NVCC. On Windows, the compilation requires Microsoft Visual Studio. We recommend installing [Visual Studio Community Edition](https://visualstudio.microsoft.com/vs/) and adding it into `PATH` using `"C:\Program Files (x86)\Microsoft Visual Studio\<VERSION>\Community\VC\Auxiliary\Build\vcvars64.bat"`.
+
+See [Troubleshooting](./docs/troubleshooting.md) for help on common installation and run-time problems.
+
+## Getting started
+
+Pre-trained networks are stored as `*.pkl` files that can be referenced using local filenames or URLs:
+
+```bash
+# Generate an image using pre-trained AFHQv2 model ("Ours" in Figure 1, left).
+python gen_images.py --outdir=out --trunc=1 --seeds=2 \
+    --network=https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-r-afhqv2-512x512.pkl
+
+# Render a 4x2 grid of interpolations for seeds 0 through 31.
+python gen_video.py --output=lerp.mp4 --trunc=1 --seeds=0-31 --grid=4x2 \
+    --network=https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-r-afhqv2-512x512.pkl
+```
+
+Outputs from the above commands are placed under `out/*.png`, controlled by `--outdir`. Downloaded network pickles are cached under `$HOME/.cache/dnnlib`, which can be overridden by setting the `DNNLIB_CACHE_DIR` environment variable. The default PyTorch extension build directory is `$HOME/.cache/torch_extensions`, which can be overridden by setting `TORCH_EXTENSIONS_DIR`.
+
+**Docker**: You can run the above curated image example using Docker as follows:
+
+```bash
+# Build the stylegan3:latest image
+docker build --tag stylegan3 .
+
+# Run the gen_images.py script using Docker:
+docker run --gpus all -it --rm --user $(id -u):$(id -g) \
+    -v `pwd`:/scratch --workdir /scratch -e HOME=/scratch \
+    stylegan3 \
+    python gen_images.py --outdir=out --trunc=1 --seeds=2 \
+         --network=https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-r-afhqv2-512x512.pkl
+```
+
+Note: The Docker image requires NVIDIA driver release `r470` or later.
+
+The `docker run` invocation may look daunting, so let's unpack its contents here:
+
+- `--gpus all -it --rm --user $(id -u):$(id -g)`: with all GPUs enabled, run an interactive session with current user's UID/GID to avoid Docker writing files as root.
+- ``-v `pwd`:/scratch --workdir /scratch``: mount current running dir (e.g., the top of this git repo on your host machine) to `/scratch` in the container and use that as the current working dir.
+- `-e HOME=/scratch`: let PyTorch and StyleGAN3 code know where to cache temporary files such as pre-trained models and custom PyTorch extension build results. Note: if you want more fine-grained control, you can instead set `TORCH_EXTENSIONS_DIR` (for custom extensions build dir) and `DNNLIB_CACHE_DIR` (for pre-trained model download cache). You want these cache dirs to reside on persistent volumes so that their contents are retained across multiple `docker run` invocations.
+
+## Interactive visualization
+
+This release contains an interactive model visualization tool that can be used to explore various characteristics of a trained model.  To start it, run:
+
+```bash
+python visualizer.py
+```
+
+<a href="./docs/visualizer_screen0.png"><img alt="Visualizer screenshot" src="./docs/visualizer_screen0_half.png"></img></a>
+
+## Using networks from Python
+
+You can use pre-trained networks in your own Python code as follows:
+
+```python
+with open('ffhq.pkl', 'rb') as f:
+    G = pickle.load(f)['G_ema'].cuda()  # torch.nn.Module
+z = torch.randn([1, G.z_dim]).cuda()    # latent codes
+c = None                                # class labels (not used in this example)
+img = G(z, c)                           # NCHW, float32, dynamic range [-1, +1], no truncation
+```
+
+The above code requires `torch_utils` and `dnnlib` to be accessible via `PYTHONPATH`. It does not need source code for the networks themselves &mdash; their class definitions are loaded from the pickle via `torch_utils.persistence`.
+
+The pickle contains three networks. `'G'` and `'D'` are instantaneous snapshots taken during training, and `'G_ema'` represents a moving average of the generator weights over several training steps. The networks are regular instances of `torch.nn.Module`, with all of their parameters and buffers placed on the CPU at import and gradient computation disabled by default.
+
+The generator consists of two submodules, `G.mapping` and `G.synthesis`, that can be executed separately. They also support various additional options:
+
+```python
+w = G.mapping(z, c, truncation_psi=0.5, truncation_cutoff=8)
+img = G.synthesis(w, noise_mode='const', force_fp32=True)
+```
+
+Please refer to [`gen_images.py`](./gen_images.py) for complete code example.
+
+## Preparing datasets
+
+Datasets are stored as uncompressed ZIP archives containing uncompressed PNG files and a metadata file `dataset.json` for labels. Custom datasets can be created from a folder containing images; see [`python dataset_tool.py --help`](./docs/dataset-tool-help.txt) for more information. Alternatively, the folder can also be used directly as a dataset, without running it through `dataset_tool.py` first, but doing so may lead to suboptimal performance.
+
+**FFHQ**: Download the [Flickr-Faces-HQ dataset](https://github.com/NVlabs/ffhq-dataset) as 1024x1024 images and create a zip archive using `dataset_tool.py`:
+
+```bash
+# Original 1024x1024 resolution.
+python dataset_tool.py --source=/tmp/images1024x1024 --dest=~/datasets/ffhq-1024x1024.zip
+
+# Scaled down 256x256 resolution.
+python dataset_tool.py --source=/tmp/images1024x1024 --dest=~/datasets/ffhq-256x256.zip \
+    --resolution=256x256
+```
+
+See the [FFHQ README](https://github.com/NVlabs/ffhq-dataset) for information on how to obtain the unaligned FFHQ dataset images. Use the same steps as above to create a ZIP archive for training and validation.
+
+**MetFaces**: Download the [MetFaces dataset](https://github.com/NVlabs/metfaces-dataset) and create a ZIP archive:
+
+```bash
+python dataset_tool.py --source=~/downloads/metfaces/images --dest=~/datasets/metfaces-1024x1024.zip
+```
+
+See the [MetFaces README](https://github.com/NVlabs/metfaces-dataset) for information on how to obtain the unaligned MetFaces dataset images. Use the same steps as above to create a ZIP archive for training and validation.
+
+**AFHQv2**: Download the [AFHQv2 dataset](https://github.com/clovaai/stargan-v2/blob/master/README.md#animal-faces-hq-dataset-afhq) and create a ZIP archive:
+
+```bash
+python dataset_tool.py --source=~/downloads/afhqv2 --dest=~/datasets/afhqv2-512x512.zip
+```
+
+Note that the above command creates a single combined dataset using all images of all three classes (cats, dogs, and wild animals), matching the setup used in the StyleGAN3 paper. Alternatively, you can also create a separate dataset for each class:
+
+```bash
+python dataset_tool.py --source=~/downloads/afhqv2/train/cat --dest=~/datasets/afhqv2cat-512x512.zip
+python dataset_tool.py --source=~/downloads/afhqv2/train/dog --dest=~/datasets/afhqv2dog-512x512.zip
+python dataset_tool.py --source=~/downloads/afhqv2/train/wild --dest=~/datasets/afhqv2wild-512x512.zip
+```
+
+## Training
+
+You can train new networks using `train.py`. For example:
+
+```bash
+# Train StyleGAN3-T for AFHQv2 using 8 GPUs.
+python train.py --outdir=~/training-runs --cfg=stylegan3-t --data=~/datasets/afhqv2-512x512.zip \
+    --gpus=8 --batch=32 --gamma=8.2 --mirror=1
+
+# Fine-tune StyleGAN3-R for MetFaces-U using 1 GPU, starting from the pre-trained FFHQ-U pickle.
+python train.py --outdir=~/training-runs --cfg=stylegan3-r --data=~/datasets/metfacesu-1024x1024.zip \
+    --gpus=8 --batch=32 --gamma=6.6 --mirror=1 --kimg=5000 --snap=5 \
+    --resume=https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-r-ffhqu-1024x1024.pkl
+
+# Train StyleGAN2 for FFHQ at 1024x1024 resolution using 8 GPUs.
+python train.py --outdir=~/training-runs --cfg=stylegan2 --data=~/datasets/ffhq-1024x1024.zip \
+    --gpus=8 --batch=32 --gamma=10 --mirror=1 --aug=noaug
+```
+
+Note that the result quality and training time depend heavily on the exact set of options. The most important ones (`--gpus`, `--batch`, and `--gamma`) must be specified explicitly, and they should be selected with care. See [`python train.py --help`](./docs/train-help.txt) for the full list of options and [Training configurations](./docs/configs.md) for general guidelines &amp; recommendations, along with the expected training speed &amp; memory usage in different scenarios.
+
+The results of each training run are saved to a newly created directory, for example `~/training-runs/00000-stylegan3-t-afhqv2-512x512-gpus8-batch32-gamma8.2`. The training loop exports network pickles (`network-snapshot-<KIMG>.pkl`) and random image grids (`fakes<KIMG>.png`) at regular intervals (controlled by `--snap`). For each exported pickle, it evaluates FID (controlled by `--metrics`) and logs the result in `metric-fid50k_full.jsonl`. It also records various statistics in `training_stats.jsonl`, as well as `*.tfevents` if TensorBoard is installed.
+
+## Quality metrics
+
+By default, `train.py` automatically computes FID for each network pickle exported during training. We recommend inspecting `metric-fid50k_full.jsonl` (or TensorBoard) at regular intervals to monitor the training progress. When desired, the automatic computation can be disabled with `--metrics=none` to speed up the training slightly.
+
+Additional quality metrics can also be computed after the training:
+
+```bash
+# Previous training run: look up options automatically, save result to JSONL file.
+python calc_metrics.py --metrics=eqt50k_int,eqr50k \
+    --network=~/training-runs/00000-stylegan3-r-mydataset/network-snapshot-000000.pkl
+
+# Pre-trained network pickle: specify dataset explicitly, print result to stdout.
+python calc_metrics.py --metrics=fid50k_full --data=~/datasets/ffhq-1024x1024.zip --mirror=1 \
+    --network=https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-t-ffhq-1024x1024.pkl
+```
+
+The first example looks up the training configuration and performs the same operation as if `--metrics=eqt50k_int,eqr50k` had been specified during training. The second example downloads a pre-trained network pickle, in which case the values of `--data` and `--mirror` must be specified explicitly.
+
+Note that the metrics can be quite expensive to compute (up to 1h), and many of them have an additional one-off cost for each new dataset (up to 30min). Also note that the evaluation is done using a different random seed each time, so the results will vary if the same metric is computed multiple times.
+
+Recommended metrics:
+* `fid50k_full`: Fr&eacute;chet inception distance<sup>[1]</sup> against the full dataset.
+* `kid50k_full`: Kernel inception distance<sup>[2]</sup> against the full dataset.
+* `pr50k3_full`: Precision and recall<sup>[3]</sup> againt the full dataset.
+* `ppl2_wend`: Perceptual path length<sup>[4]</sup> in W, endpoints, full image.
+* `eqt50k_int`: Equivariance<sup>[5]</sup> w.r.t. integer translation (EQ-T).
+* `eqt50k_frac`: Equivariance w.r.t. fractional translation (EQ-T<sub>frac</sub>).
+* `eqr50k`: Equivariance w.r.t. rotation (EQ-R).
+
+Legacy metrics:
+* `fid50k`: Fr&eacute;chet inception distance against 50k real images.
+* `kid50k`: Kernel inception distance against 50k real images.
+* `pr50k3`: Precision and recall against 50k real images.
+* `is50k`: Inception score<sup>[6]</sup> for CIFAR-10.
+
+References:
+1. [GANs Trained by a Two Time-Scale Update Rule Converge to a Local Nash Equilibrium](https://arxiv.org/abs/1706.08500), Heusel et al. 2017
+2. [Demystifying MMD GANs](https://arxiv.org/abs/1801.01401), Bi&nacute;kowski et al. 2018
+3. [Improved Precision and Recall Metric for Assessing Generative Models](https://arxiv.org/abs/1904.06991), Kynk&auml;&auml;nniemi et al. 2019
+4. [A Style-Based Generator Architecture for Generative Adversarial Networks](https://arxiv.org/abs/1812.04948), Karras et al. 2018
+5. [Alias-Free Generative Adversarial Networks](https://nvlabs.github.io/stylegan3), Karras et al. 2021
+6. [Improved Techniques for Training GANs](https://arxiv.org/abs/1606.03498), Salimans et al. 2016
+
+## Spectral analysis
+
+The easiest way to inspect the spectral properties of a given generator is to use the built-in FFT mode in `visualizer.py`. In addition, you can visualize average 2D power spectra (Appendix A, Figure 15) as follows:
+
+```bash
+# Calculate dataset mean and std, needed in subsequent steps.
+python avg_spectra.py stats --source=~/datasets/ffhq-1024x1024.zip
+
+# Calculate average spectrum for the training data.
+python avg_spectra.py calc --source=~/datasets/ffhq-1024x1024.zip \
+    --dest=tmp/training-data.npz --mean=112.684 --std=69.509
+
+# Calculate average spectrum for a pre-trained generator.
+python avg_spectra.py calc \
+    --source=https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-r-ffhq-1024x1024.pkl \
+    --dest=tmp/stylegan3-r.npz --mean=112.684 --std=69.509 --num=70000
+
+# Display results.
+python avg_spectra.py heatmap tmp/training-data.npz
+python avg_spectra.py heatmap tmp/stylegan3-r.npz
+python avg_spectra.py slices tmp/training-data.npz tmp/stylegan3-r.npz
+```
+
+<a href="./docs/avg_spectra_screen0.png"><img alt="Average spectra screenshot" src="./docs/avg_spectra_screen0_half.png"></img></a>
+
+## License
+
+Copyright &copy; 2021, NVIDIA Corporation & affiliates. All rights reserved.
+
+This work is made available under the [Nvidia Source Code License](https://github.com/NVlabs/stylegan3/blob/main/LICENSE.txt).
+
+## Citation
+
+```
+@inproceedings{Karras2021,
+  author = {Tero Karras and Miika Aittala and Samuli Laine and Erik H\"ark\"onen and Janne Hellsten and Jaakko Lehtinen and Timo Aila},
+  title = {Alias-Free Generative Adversarial Networks},
+  booktitle = {Proc. NeurIPS},
+  year = {2021}
+}
+```
+
+## Development
+
+This is a research reference implementation and is treated as a one-time code drop. As such, we do not accept outside code contributions in the form of pull requests.
+
+## Acknowledgements
+
+We thank David Luebke, Ming-Yu Liu, Koki Nagano, Tuomas Kynk&auml;&auml;nniemi, and Timo Viitanen for reviewing early drafts and helpful suggestions. Fr&eacute;do Durand for early discussions. Tero Kuosmanen for maintaining our compute infrastructure. AFHQ authors for an updated version of their dataset. Getty Images for the training images in the Beaches dataset. We did not receive external funding or additional revenues for this project.

stylegan3-fun/avg_spectra.py

@@ -0,0 +1,276 @@
+# 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.
+
+"""Compare average power spectra between real and generated images,
+or between multiple generators."""
+
+import os
+import numpy as np
+import torch
+import torch.fft
+import scipy.ndimage
+import matplotlib.pyplot as plt
+import click
+import tqdm
+import dnnlib
+
+import legacy
+from training import dataset
+
+#----------------------------------------------------------------------------
+# Setup an iterator for streaming images, in uint8 NCHW format, based on the
+# respective command line options.
+
+def stream_source_images(source, num, seed, device, data_loader_kwargs=None): # => num_images, image_size, image_iter
+    ext = source.split('.')[-1].lower()
+    if data_loader_kwargs is None:
+        data_loader_kwargs = dict(pin_memory=True, num_workers=3, prefetch_factor=2)
+
+    if ext == 'pkl':
+        if num is None:
+            raise click.ClickException('--num is required when --source points to network pickle')
+        with dnnlib.util.open_url(source) as f:
+            G = legacy.load_network_pkl(f)['G_ema'].to(device)
+        def generate_image(seed):
+            rnd = np.random.RandomState(seed)
+            z = torch.from_numpy(rnd.randn(1, G.z_dim)).to(device)
+            c = torch.zeros([1, G.c_dim], device=device)
+            if G.c_dim > 0:
+                c[:, rnd.randint(G.c_dim)] = 1
+            return (G(z=z, c=c) * 127.5 + 128).clamp(0, 255).to(torch.uint8)
+        _ = generate_image(seed) # warm up
+        image_iter = (generate_image(seed + idx) for idx in range(num))
+        return num, G.img_resolution, image_iter
+
+    elif ext == 'zip' or os.path.isdir(source):
+        dataset_obj = dataset.ImageFolderDataset(path=source, max_size=num, random_seed=seed)
+        if num is not None and num != len(dataset_obj):
+            raise click.ClickException(f'--source contains fewer than {num} images')
+        data_loader = torch.utils.data.DataLoader(dataset_obj, batch_size=1, **data_loader_kwargs)
+        image_iter = (image.to(device) for image, _label in data_loader)
+        return len(dataset_obj), dataset_obj.resolution, image_iter
+
+    else:
+        raise click.ClickException('--source must point to network pickle, dataset zip, or directory')
+
+#----------------------------------------------------------------------------
+# Load average power spectrum from the specified .npz file and construct
+# the corresponding heatmap for visualization.
+
+def construct_heatmap(npz_file, smooth):
+    npz_data = np.load(npz_file)
+    spectrum = npz_data['spectrum']
+    image_size = npz_data['image_size']
+    hmap = np.log10(spectrum) * 10 # dB
+    hmap = np.fft.fftshift(hmap)
+    hmap = np.concatenate([hmap, hmap[:1, :]], axis=0)
+    hmap = np.concatenate([hmap, hmap[:, :1]], axis=1)
+    if smooth > 0:
+        sigma = spectrum.shape[0] / image_size * smooth
+        hmap = scipy.ndimage.gaussian_filter(hmap, sigma=sigma, mode='nearest')
+    return hmap, image_size
+
+#----------------------------------------------------------------------------
+
+@click.group()
+def main():
+    """Compare average power spectra between real and generated images,
+    or between multiple generators.
+
+    Example:
+
+    \b
+    # Calculate dataset mean and std, needed in subsequent steps.
+    python avg_spectra.py stats --source=~/datasets/ffhq-1024x1024.zip
+
+    \b
+    # Calculate average spectrum for the training data.
+    python avg_spectra.py calc --source=~/datasets/ffhq-1024x1024.zip \\
+        --dest=tmp/training-data.npz --mean=112.684 --std=69.509
+
+    \b
+    # Calculate average spectrum for a pre-trained generator.
+    python avg_spectra.py calc \\
+        --source=https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-r-ffhq-1024x1024.pkl \\
+        --dest=tmp/stylegan3-r.npz --mean=112.684 --std=69.509 --num=70000
+
+    \b
+    # Display results.
+    python avg_spectra.py heatmap tmp/training-data.npz
+    python avg_spectra.py heatmap tmp/stylegan3-r.npz
+    python avg_spectra.py slices tmp/training-data.npz tmp/stylegan3-r.npz
+
+    \b
+    # Save as PNG.
+    python avg_spectra.py heatmap tmp/training-data.npz --save=tmp/training-data.png --dpi=300
+    python avg_spectra.py heatmap tmp/stylegan3-r.npz --save=tmp/stylegan3-r.png --dpi=300
+    python avg_spectra.py slices tmp/training-data.npz tmp/stylegan3-r.npz --save=tmp/slices.png --dpi=300
+    """
+
+#----------------------------------------------------------------------------
+
+@main.command()
+@click.option('--source', help='Network pkl, dataset zip, or directory', metavar='[PKL|ZIP|DIR]', required=True)
+@click.option('--num', help='Number of images to process  [default: all]', metavar='INT', type=click.IntRange(min=1))
+@click.option('--seed', help='Random seed for selecting the images', metavar='INT', type=click.IntRange(min=0), default=0, show_default=True)
+def stats(source, num, seed, device=torch.device('cuda')):
+    """Calculate dataset mean and standard deviation needed by 'calc'."""
+    torch.multiprocessing.set_start_method('spawn')
+    num_images, _image_size, image_iter = stream_source_images(source=source, num=num, seed=seed, device=device)
+
+    # Accumulate moments.
+    moments = torch.zeros([3], dtype=torch.float64, device=device)
+    for image in tqdm.tqdm(image_iter, total=num_images):
+        image = image.to(torch.float64)
+        moments += torch.stack([torch.ones_like(image).sum(), image.sum(), image.square().sum()])
+    moments = moments / moments[0]
+
+    # Compute mean and standard deviation.
+    mean = moments[1]
+    std = (moments[2] - moments[1].square()).sqrt()
+    print(f'--mean={mean:g} --std={std:g}')
+
+#----------------------------------------------------------------------------
+
+@main.command()
+@click.option('--source', help='Network pkl, dataset zip, or directory', metavar='[PKL|ZIP|DIR]', required=True)
+@click.option('--dest', help='Where to store the result', metavar='NPZ', required=True)
+@click.option('--mean', help='Dataset mean for whitening', metavar='FLOAT', type=float, required=True)
+@click.option('--std', help='Dataset standard deviation for whitening', metavar='FLOAT', type=click.FloatRange(min=0), required=True)
+@click.option('--num', help='Number of images to process  [default: all]', metavar='INT', type=click.IntRange(min=1))
+@click.option('--seed', help='Random seed for selecting the images', metavar='INT', type=click.IntRange(min=0), default=0, show_default=True)
+@click.option('--beta', help='Shape parameter for the Kaiser window', metavar='FLOAT', type=click.FloatRange(min=0), default=8, show_default=True)
+@click.option('--interp', help='Frequency-domain interpolation factor', metavar='INT', type=click.IntRange(min=1), default=4, show_default=True)
+def calc(source, dest, mean, std, num, seed, beta, interp, device=torch.device('cuda')):
+    """Calculate average power spectrum and store it in .npz file."""
+    torch.multiprocessing.set_start_method('spawn')
+    num_images, image_size, image_iter = stream_source_images(source=source, num=num, seed=seed, device=device)
+    spectrum_size = image_size * interp
+    padding = spectrum_size - image_size
+
+    # Setup window function.
+    window = torch.kaiser_window(image_size, periodic=False, beta=beta, device=device)
+    window *= window.square().sum().rsqrt()
+    window = window.ger(window).unsqueeze(0).unsqueeze(1)
+
+    # Accumulate power spectrum.
+    spectrum = torch.zeros([spectrum_size, spectrum_size], dtype=torch.float64, device=device)
+    for image in tqdm.tqdm(image_iter, total=num_images):
+        image = (image.to(torch.float64) - mean) / std
+        image = torch.nn.functional.pad(image * window, [0, padding, 0, padding])
+        spectrum += torch.fft.fftn(image, dim=[2,3]).abs().square().mean(dim=[0,1])
+    spectrum /= num_images
+
+    # Save result.
+    if os.path.dirname(dest):
+        os.makedirs(os.path.dirname(dest), exist_ok=True)
+    np.savez(dest, spectrum=spectrum.cpu().numpy(), image_size=image_size)
+
+#----------------------------------------------------------------------------
+
+@main.command()
+@click.argument('npz-file', nargs=1)
+@click.option('--save', help='Save the plot and exit', metavar='[PNG|PDF|...]')
+@click.option('--dpi', help='Figure resolution', metavar='FLOAT', type=click.FloatRange(min=1), default=100, show_default=True)
+@click.option('--smooth', help='Amount of smoothing', metavar='FLOAT', type=click.FloatRange(min=0), default=1.25, show_default=True)
+def heatmap(npz_file, save, smooth, dpi):
+    """Visualize 2D heatmap based on the given .npz file."""
+    hmap, image_size = construct_heatmap(npz_file=npz_file, smooth=smooth)
+
+    # Setup plot.
+    plt.figure(figsize=[6, 4.8], dpi=dpi, tight_layout=True)
+    freqs = np.linspace(-0.5, 0.5, num=hmap.shape[0], endpoint=True) * image_size
+    ticks = np.linspace(freqs[0], freqs[-1], num=5, endpoint=True)
+    levels = np.linspace(-40, 20, num=13, endpoint=True)
+
+    # Draw heatmap.
+    plt.xlim(ticks[0], ticks[-1])
+    plt.ylim(ticks[0], ticks[-1])
+    plt.xticks(ticks)
+    plt.yticks(ticks)
+    plt.contourf(freqs, freqs, hmap, levels=levels, extend='both', cmap='Blues')
+    plt.gca().set_aspect('equal')
+    plt.colorbar(ticks=levels)
+    plt.contour(freqs, freqs, hmap, levels=levels, extend='both', linestyles='solid', linewidths=1, colors='midnightblue', alpha=0.2)
+
+    # Display or save.
+    if save is None:
+        plt.show()
+    else:
+        if os.path.dirname(save):
+            os.makedirs(os.path.dirname(save), exist_ok=True)
+        plt.savefig(save)
+
+#----------------------------------------------------------------------------
+
+@main.command()
+@click.argument('npz-files', nargs=-1, required=True)
+@click.option('--save', help='Save the plot and exit', metavar='[PNG|PDF|...]')
+@click.option('--dpi', help='Figure resolution', metavar='FLOAT', type=click.FloatRange(min=1), default=100, show_default=True)
+@click.option('--smooth', help='Amount of smoothing', metavar='FLOAT', type=click.FloatRange(min=0), default=0, show_default=True)
+def slices(npz_files, save, dpi, smooth):
+    """Visualize 1D slices based on the given .npz files."""
+    cases = [dnnlib.EasyDict(npz_file=npz_file) for npz_file in npz_files]
+    for c in cases:
+        c.hmap, c.image_size = construct_heatmap(npz_file=c.npz_file, smooth=smooth)
+        c.label = os.path.splitext(os.path.basename(c.npz_file))[0]
+
+    # Check consistency.
+    image_size = cases[0].image_size
+    hmap_size = cases[0].hmap.shape[0]
+    if any(c.image_size != image_size or c.hmap.shape[0] != hmap_size for c in cases):
+        raise click.ClickException('All .npz must have the same resolution')
+
+    # Setup plot.
+    plt.figure(figsize=[12, 4.6], dpi=dpi, tight_layout=True)
+    hmap_center = hmap_size // 2
+    hmap_range = np.arange(hmap_center, hmap_size)
+    freqs0 = np.linspace(0, image_size / 2, num=(hmap_size // 2 + 1), endpoint=True)
+    freqs45 = np.linspace(0, image_size / np.sqrt(2), num=(hmap_size // 2 + 1), endpoint=True)
+    xticks0 = np.linspace(freqs0[0], freqs0[-1], num=9, endpoint=True)
+    xticks45 = np.round(np.linspace(freqs45[0], freqs45[-1], num=9, endpoint=True))
+    yticks = np.linspace(-50, 30, num=9, endpoint=True)
+
+    # Draw 0 degree slice.
+    plt.subplot(1, 2, 1)
+    plt.title('0\u00b0 slice')
+    plt.xlim(xticks0[0], xticks0[-1])
+    plt.ylim(yticks[0], yticks[-1])
+    plt.xticks(xticks0)
+    plt.yticks(yticks)
+    for c in cases:
+        plt.plot(freqs0, c.hmap[hmap_center, hmap_range], label=c.label)
+    plt.grid()
+    plt.legend(loc='upper right')
+
+    # Draw 45 degree slice.
+    plt.subplot(1, 2, 2)
+    plt.title('45\u00b0 slice')
+    plt.xlim(xticks45[0], xticks45[-1])
+    plt.ylim(yticks[0], yticks[-1])
+    plt.xticks(xticks45)
+    plt.yticks(yticks)
+    for c in cases:
+        plt.plot(freqs45, c.hmap[hmap_range, hmap_range], label=c.label)
+    plt.grid()
+    plt.legend(loc='upper right')
+
+    # Display or save.
+    if save is None:
+        plt.show()
+    else:
+        if os.path.dirname(save):
+            os.makedirs(os.path.dirname(save), exist_ok=True)
+        plt.savefig(save)
+
+#----------------------------------------------------------------------------
+
+if __name__ == "__main__":
+    main() # pylint: disable=no-value-for-parameter
+
+#----------------------------------------------------------------------------

stylegan3-fun/calc_metrics.py

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

stylegan3-fun/dataset_tool.py

@@ -0,0 +1,547 @@
+# 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.
+
+"""Tool for creating ZIP/PNG based datasets."""
+
+import functools
+import gzip
+import io
+import json
+import os
+import pickle
+import re
+import sys
+import tarfile
+import zipfile
+from pathlib import Path
+from typing import Callable, Optional, Tuple, Union
+
+import click
+import numpy as np
+import PIL.Image
+from tqdm import tqdm
+from torch_utils import gen_utils
+
+
+# ----------------------------------------------------------------------------
+
+
+def error(msg):
+    print('Error: ' + msg)
+    sys.exit(1)
+
+
+# ----------------------------------------------------------------------------
+
+
+def parse_tuple(s: str) -> Tuple[int, int]:
+    """Parse a 'M,N' or 'MxN' integer tuple.
+
+    Example:
+        '4x2' returns (4,2)
+        '0,1' returns (0,1)
+    """
+    m = re.match(r'^(\d+)[x,](\d+)$', s)
+    if m:
+        return int(m.group(1)), int(m.group(2))
+    raise ValueError(f'cannot parse tuple {s}')
+
+
+# ----------------------------------------------------------------------------
+
+
+def maybe_min(a: int, b: Optional[int]) -> int:
+    if b is not None:
+        return min(a, b)
+    return a
+
+
+# ----------------------------------------------------------------------------
+
+
+def file_ext(name: Union[str, Path]) -> str:
+    return str(name).split('.')[-1]
+
+
+# ----------------------------------------------------------------------------
+
+
+def is_image_ext(fname: Union[str, Path]) -> bool:
+    ext = file_ext(fname).lower()
+    return f'.{ext}' in PIL.Image.EXTENSION  # type: ignore
+
+
+# ----------------------------------------------------------------------------
+
+
+def open_image_folder(source_dir, force_channels: int = None, *, max_images: Optional[int], subfolders_as_labels: Optional[bool] = False):
+    input_images = [str(f) for f in sorted(Path(source_dir).rglob('*')) if is_image_ext(f) and os.path.isfile(f)]
+
+    # Load labels.
+    labels = {}
+    meta_fname = os.path.join(source_dir, 'dataset.json')
+    if os.path.isfile(meta_fname) and not subfolders_as_labels:
+        # The `dataset.json` file exists and will be used (whether empty or not)
+        with open(meta_fname, 'r') as file:
+            labels = json.load(file)['labels']
+            if labels is not None:
+                labels = {x[0]: x[1] for x in labels}
+            else:
+                labels = {}
+    elif subfolders_as_labels:
+        # Use the folders in the directory as the labels themselves
+        # Get the subfolder names from the input_images names
+        labels = {os.path.relpath(fname, source_dir).replace('\\', '/'): os.path.basename(os.path.dirname(fname)) for fname in input_images}
+        # Change folder name (value) to a number (from 0 to n-1)
+        label_names = list(set(labels.values()))
+        label_names.sort()
+        labels = {fname: label_names.index(label) for fname, label in labels.items()}
+        print(f'Conditional dataset has {len(label_names)} labels! Saving to `class_labels.txt` in the source directory...')
+        with open(os.path.join(source_dir, 'class_labels.txt'), 'w') as f:
+            # Write, one per line, the index and the label name
+            for i, label in enumerate(label_names):
+                f.write(f'{i}: {label}\n')
+
+    max_idx = maybe_min(len(input_images), max_images)
+
+    def iterate_images():
+        for idx, fname in enumerate(input_images):
+            arch_fname = os.path.relpath(fname, source_dir)
+            arch_fname = arch_fname.replace('\\', '/')
+            # Adding Pull #39 from Andreas Jansson: https://github.com/NVlabs/stylegan3/pull/39
+            try:
+                img = PIL.Image.open(fname)  # Let PIL handle the mode
+                # Convert grayscale image to RGB
+                if img.mode == 'L':
+                    img = img.convert('RGB')
+                # Force the number of channels if so requested
+                if force_channels is not None:
+                    img = img.convert(gen_utils.channels_dict[int(force_channels)])
+                img = np.array(img)
+            except Exception as e:
+                sys.stderr.write(f'Failed to read {fname}: {e}')
+                continue
+            yield dict(img=img, label=labels.get(arch_fname))
+            if idx >= max_idx-1:
+                break
+    return max_idx, iterate_images()
+
+
+# ----------------------------------------------------------------------------
+
+
+def open_image_zip(source, force_channels: int = None, *, max_images: Optional[int]):
+    with zipfile.ZipFile(source, mode='r') as z:
+        input_images = [str(f) for f in sorted(z.namelist()) if is_image_ext(f)]
+
+        # Load labels.
+        labels = {}
+        if 'dataset.json' in z.namelist():
+            with z.open('dataset.json', 'r') as file:
+                labels = json.load(file)['labels']
+                if labels is not None:
+                    labels = {x[0]: x[1] for x in labels}
+                else:
+                    labels = {}
+
+    max_idx = maybe_min(len(input_images), max_images)
+
+    def iterate_images():
+        with zipfile.ZipFile(source, mode='r') as z:
+            for idx, fname in enumerate(input_images):
+                with z.open(fname, 'r') as file:
+                    # Same as above: PR #39 by Andreas Jansson and turn Grayscale to RGB
+                    try:
+                        img = PIL.Image.open(file) # type: ignore
+                        if img.mode == 'L':
+                            img = img.convert('RGB')
+                        # Force the number of channels if so requested
+                        if force_channels is not None:
+                            img = img.convert(gen_utils.channels_dict[int(force_channels)])
+                        img = np.array(img)
+                    except Exception as e:
+                        sys.stderr.write(f'Failed to read {fname}: {e}')
+                        continue
+                yield dict(img=img, label=labels.get(fname))
+                if idx >= max_idx-1:
+                    break
+    return max_idx, iterate_images()
+
+
+# ----------------------------------------------------------------------------
+
+
+def open_lmdb(lmdb_dir: str, *, max_images: Optional[int]):
+    import cv2  # pip install opencv-python # pylint: disable=import-error
+    import lmdb  # pip install lmdb # pylint: disable=import-error
+
+    with lmdb.open(lmdb_dir, readonly=True, lock=False).begin(write=False) as txn:
+        max_idx = maybe_min(txn.stat()['entries'], max_images)
+
+    def iterate_images():
+        with lmdb.open(lmdb_dir, readonly=True, lock=False).begin(write=False) as txn:
+            for idx, (_key, value) in enumerate(txn.cursor()):
+                try:
+                    try:
+                        img = cv2.imdecode(np.frombuffer(value, dtype=np.uint8), 1)
+                        if img is None:
+                            raise IOError('cv2.imdecode failed')
+                        img = img[:, :, ::-1]  # BGR => RGB
+                    except IOError:
+                        img = np.array(PIL.Image.open(io.BytesIO(value)))
+                    yield dict(img=img, label=None)
+                    if idx >= max_idx-1:
+                        break
+                except:
+                    print(sys.exc_info()[1])
+
+    return max_idx, iterate_images()
+
+
+# ----------------------------------------------------------------------------
+
+
+def open_cifar10(tarball: str, *, max_images: Optional[int]):
+    images = []
+    labels = []
+
+    with tarfile.open(tarball, 'r:gz') as tar:
+        for batch in range(1, 6):
+            member = tar.getmember(f'cifar-10-batches-py/data_batch_{batch}')
+            with tar.extractfile(member) as file:
+                data = pickle.load(file, encoding='latin1')
+            images.append(data['data'].reshape(-1, 3, 32, 32))
+            labels.append(data['labels'])
+
+    images = np.concatenate(images)
+    labels = np.concatenate(labels)
+    images = images.transpose([0, 2, 3, 1])  # NCHW -> NHWC
+    assert images.shape == (50000, 32, 32, 3) and images.dtype == np.uint8
+    assert labels.shape == (50000,) and labels.dtype in [np.int32, np.int64]
+    assert np.min(images) == 0 and np.max(images) == 255
+    assert np.min(labels) == 0 and np.max(labels) == 9
+
+    max_idx = maybe_min(len(images), max_images)
+
+    def iterate_images():
+        for idx, img in enumerate(images):
+            yield dict(img=img, label=int(labels[idx]))
+            if idx >= max_idx-1:
+                break
+
+    return max_idx, iterate_images()
+
+
+# ----------------------------------------------------------------------------
+
+
+def open_mnist(images_gz: str, *, max_images: Optional[int]):
+    labels_gz = images_gz.replace('-images-idx3-ubyte.gz', '-labels-idx1-ubyte.gz')
+    assert labels_gz != images_gz
+    images = []
+    labels = []
+
+    with gzip.open(images_gz, 'rb') as f:
+        images = np.frombuffer(f.read(), np.uint8, offset=16)
+    with gzip.open(labels_gz, 'rb') as f:
+        labels = np.frombuffer(f.read(), np.uint8, offset=8)
+
+    images = images.reshape(-1, 28, 28)
+    images = np.pad(images, [(0, 0), (2, 2), (2, 2)], 'constant', constant_values=0)
+    assert images.shape == (60000, 32, 32) and images.dtype == np.uint8
+    assert labels.shape == (60000,) and labels.dtype == np.uint8
+    assert np.min(images) == 0 and np.max(images) == 255
+    assert np.min(labels) == 0 and np.max(labels) == 9
+
+    max_idx = maybe_min(len(images), max_images)
+
+    def iterate_images():
+        for idx, img in enumerate(images):
+            yield dict(img=img, label=int(labels[idx]))
+            if idx >= max_idx-1:
+                break
+
+    return max_idx, iterate_images()
+
+
+# ----------------------------------------------------------------------------
+
+
+def make_transform(
+    transform: Optional[str],
+    output_width: Optional[int],
+    output_height: Optional[int]
+) -> Callable[[np.ndarray], Optional[np.ndarray]]:
+    def scale(width, height, img):
+        w = img.shape[1]
+        h = img.shape[0]
+        if width == w and height == h:
+            return img
+        img = PIL.Image.fromarray(img)
+        ww = width if width is not None else w
+        hh = height if height is not None else h
+        img = img.resize((ww, hh), PIL.Image.LANCZOS)
+        return np.array(img)
+
+    def center_crop(width, height, img):
+        crop = np.min(img.shape[:2])
+        img = img[(img.shape[0] - crop) // 2: (img.shape[0] + crop) // 2,
+                  (img.shape[1] - crop) // 2: (img.shape[1] + crop) // 2]
+        img = PIL.Image.fromarray(img, gen_utils.channels_dict[img.shape[2]])
+        img = img.resize((width, height), PIL.Image.LANCZOS)
+        return np.array(img)
+
+    def center_crop_wide(width, height, img):
+        ch = int(np.round(width * img.shape[0] / img.shape[1]))
+        if img.shape[1] < width or ch < height:
+            return None
+
+        img = img[(img.shape[0] - ch) // 2: (img.shape[0] + ch) // 2]
+        img = PIL.Image.fromarray(img, gen_utils.channels_dict[img.shape[2]])
+        img = img.resize((width, height), PIL.Image.LANCZOS)
+        img = np.array(img)
+
+        canvas = np.zeros([width, width, 3], dtype=np.uint8)
+        canvas[(width - height) // 2 : (width + height) // 2, :] = img
+        return canvas
+
+    def center_crop_tall(width, height, img):
+        ch = int(np.round(height * img.shape[1] / img.shape[0]))
+        if img.shape[0] < height or ch < width:
+            return None
+
+        img = img[:, (img.shape[1] - ch) // 2: (img.shape[1] + ch) // 2]  # center-crop: [width0, height0, C] -> [width0, height, C]
+        img = PIL.Image.fromarray(img, gen_utils.channels_dict[img.shape[2]])
+        img = img.resize((width, height), PIL.Image.LANCZOS)  # resize: [width0, height, 3] -> [width, height, 3]
+        img = np.array(img)
+
+        canvas = np.zeros([height, height, 3], dtype=np.uint8)  # square canvas
+        canvas[:, (height - width) // 2: (height + width) // 2] = img  # replace the middle with img
+        return canvas
+
+    if transform is None:
+        return functools.partial(scale, output_width, output_height)
+    if transform == 'center-crop':
+        if (output_width is None) or (output_height is None):
+            error(f'must specify --resolution=WxH when using {transform} transform')
+        return functools.partial(center_crop, output_width, output_height)
+    if transform == 'center-crop-wide':
+        if (output_width is None) or (output_height is None):
+            error(f'must specify --resolution=WxH when using {transform} transform')
+        return functools.partial(center_crop_wide, output_width, output_height)
+    if transform == 'center-crop-tall':
+        if (output_width is None) or (output_height is None):
+            error(f'must specify --resolution=WxH when using {transform} transform')
+        return functools.partial(center_crop_tall, output_width, output_height)
+    assert False, 'unknown transform'
+
+
+# ----------------------------------------------------------------------------
+
+
+def open_dataset(source, force_channels, *, max_images: Optional[int], subfolders_as_labels: Optional[bool] = False):
+    if os.path.isdir(source):
+        if source.rstrip('/').endswith('_lmdb'):
+            return open_lmdb(source, max_images=max_images)
+        else:
+            return open_image_folder(source, force_channels, max_images=max_images, subfolders_as_labels=subfolders_as_labels)
+    elif os.path.isfile(source):
+        if os.path.basename(source) == 'cifar-10-python.tar.gz':
+            return open_cifar10(source, max_images=max_images)
+        elif os.path.basename(source) == 'train-images-idx3-ubyte.gz':
+            return open_mnist(source, max_images=max_images)
+        elif file_ext(source) == 'zip':
+            return open_image_zip(source, force_channels, max_images=max_images)
+        else:
+            assert False, 'unknown archive type'
+    else:
+        error(f'Missing input file or directory: {source}')
+
+
+# ----------------------------------------------------------------------------
+
+
+def open_dest(dest: str) -> Tuple[str, Callable[[str, Union[bytes, str]], None], Callable[[], None]]:
+    dest_ext = file_ext(dest)
+
+    if dest_ext == 'zip':
+        if os.path.dirname(dest) != '':
+            os.makedirs(os.path.dirname(dest), exist_ok=True)
+        zf = zipfile.ZipFile(file=dest, mode='w', compression=zipfile.ZIP_STORED)
+        def zip_write_bytes(fname: str, data: Union[bytes, str]):
+            zf.writestr(fname, data)
+        return '', zip_write_bytes, zf.close
+    else:
+        # If the output folder already exists, check that it is
+        # empty.
+        #
+        # Note: creating the output directory is not strictly
+        # necessary as folder_write_bytes() also mkdirs, but it's better
+        # to give an error message earlier in case the dest folder
+        # somehow cannot be created.
+        if os.path.isdir(dest) and len(os.listdir(dest)) != 0:
+            error('--dest folder must be empty')
+        os.makedirs(dest, exist_ok=True)
+
+        def folder_write_bytes(fname: str, data: Union[bytes, str]):
+            os.makedirs(os.path.dirname(fname), exist_ok=True)
+            with open(fname, 'wb') as fout:
+                if isinstance(data, str):
+                    data = data.encode('utf8')
+                fout.write(data)
+        return dest, folder_write_bytes, lambda: None
+
+
+# ----------------------------------------------------------------------------
+
+
+@click.command()
+@click.pass_context
+@click.option('--source', help='Directory or archive name for input dataset', required=True, metavar='PATH')
+@click.option('--dest', help='Output directory or archive name for output dataset', required=True, metavar='PATH')
+@click.option('--max-images', help='Output only up to `max-images` images', type=int, default=None)
+@click.option('--force-channels', help='Force the number of channels in the image (1: grayscale, 3: RGB, 4: RGBA)', type=click.Choice(['1', '3', '4']), default=None)
+@click.option('--subfolders-as-labels', help='Use the folder names as the labels, to avoid setting up `dataset.json`', is_flag=True)
+@click.option('--transform', help='Input crop/resize mode', type=click.Choice(['center-crop', 'center-crop-wide', 'center-crop-tall']))
+@click.option('--resolution', help='Output resolution (e.g., \'512x512\')', metavar='WxH', type=parse_tuple)
+def convert_dataset(
+    ctx: click.Context,
+    source: str,
+    dest: str,
+    max_images: Optional[int],
+    force_channels: Optional[int],
+    subfolders_as_labels: Optional[bool],
+    transform: Optional[str],
+    resolution: Optional[Tuple[int, int]]
+):
+    """Convert an image dataset into a dataset archive usable with StyleGAN2 ADA PyTorch.
+
+    The input dataset format is guessed from the --source argument:
+
+    \b
+    --source *_lmdb/                    Load LSUN dataset
+    --source cifar-10-python.tar.gz     Load CIFAR-10 dataset
+    --source train-images-idx3-ubyte.gz Load MNIST dataset
+    --source path/                      Recursively load all images from path/
+    --source dataset.zip                Recursively load all images from dataset.zip
+
+    Specifying the output format and path:
+
+    \b
+    --dest /path/to/dir                 Save output files under /path/to/dir
+    --dest /path/to/dataset.zip         Save output files into /path/to/dataset.zip
+
+    The output dataset format can be either an image folder or an uncompressed zip archive.
+    Zip archives makes it easier to move datasets around file servers and clusters, and may
+    offer better training performance on network file systems.
+
+    Images within the dataset archive will be stored as uncompressed PNG.
+    Uncompresed PNGs can be efficiently decoded in the training loop.
+
+    Class labels are stored in a file called 'dataset.json' that is stored at the
+    dataset root folder.  This file has the following structure:
+
+    \b
+    {
+        "labels": [
+            ["00000/img00000000.png",6],
+            ["00000/img00000001.png",9],
+            ... repeated for every image in the dataset
+            ["00049/img00049999.png",1]
+        ]
+    }
+
+    If the 'dataset.json' file cannot be found, the dataset is interpreted as
+    not containing class labels.
+
+    Image scale/crop and resolution requirements:
+
+    Output images must be square-shaped and they must all have the same power-of-two
+    dimensions.
+
+    To scale arbitrary input image size to a specific width and height, use the
+    --resolution option.  Output resolution will be either the original
+    input resolution (if resolution was not specified) or the one specified with
+    --resolution option.
+
+    Use the --transform=center-crop or --transform=center-crop-wide options to apply a
+    center crop transform on the input image.  These options should be used with the
+    --resolution option.  For example:
+
+    \b
+    python dataset_tool.py --source LSUN/raw/cat_lmdb --dest /tmp/lsun_cat \\
+        --transform=center-crop-wide --resolution=512x384
+    """
+
+    PIL.Image.init() # type: ignore
+
+    if dest == '':
+        ctx.fail('--dest output filename or directory must not be an empty string')
+
+    num_files, input_iter = open_dataset(source, force_channels, max_images=max_images, subfolders_as_labels=subfolders_as_labels)
+    archive_root_dir, save_bytes, close_dest = open_dest(dest)
+
+    if resolution is None: resolution = (None, None)
+    transform_image = make_transform(transform, *resolution)
+
+    dataset_attrs = None
+
+    labels = []
+    for idx, image in tqdm(enumerate(input_iter), total=num_files):
+        idx_str = f'{idx:08d}'
+        archive_fname = f'{idx_str[:5]}/img{idx_str}.png'
+
+        # Apply crop and resize.
+        img = transform_image(image['img'])
+
+        # Transform may drop images.
+        if img is None:
+            continue
+
+        # Error check to require uniform image attributes across
+        # the whole dataset.
+        channels = img.shape[2] if img.ndim == 3 else 1
+        cur_image_attrs = {
+            'width': img.shape[1],
+            'height': img.shape[0],
+            'channels': channels
+        }
+        if dataset_attrs is None:
+            dataset_attrs = cur_image_attrs
+            width = dataset_attrs['width']
+            height = dataset_attrs['height']
+            if width != height:
+                error(f'Image dimensions after scale and crop are required to be square.  Got {width}x{height}')
+            if dataset_attrs['channels'] not in [1, 3, 4]:
+                error('Input images must be stored as grayscale, RGB, or RGBA')
+            if width != 2 ** int(np.floor(np.log2(width))):
+                error('Image width/height after scale and crop are required to be power-of-two')
+        elif dataset_attrs != cur_image_attrs:
+            err = [f'  dataset {k}/cur image {k}: {dataset_attrs[k]}/{cur_image_attrs[k]}' for k in dataset_attrs.keys()]  # pylint: disable=unsubscriptable-object
+            error(f'Image {archive_fname} attributes must be equal across all images of the dataset.  Got:\n' + '\n'.join(err))
+
+        # Save the image as an uncompressed PNG.
+        img = PIL.Image.fromarray(img, gen_utils.channels_dict[channels])
+        image_bits = io.BytesIO()
+        img.save(image_bits, format='png', compress_level=0, optimize=False)
+        save_bytes(os.path.join(archive_root_dir, archive_fname), image_bits.getbuffer())
+        labels.append([archive_fname, image['label']] if image['label'] is not None else None)
+
+    metadata = {
+        'labels': labels if all(x is not None for x in labels) else None
+    }
+    save_bytes(os.path.join(archive_root_dir, 'dataset.json'), json.dumps(metadata))
+    close_dest()
+
+# ----------------------------------------------------------------------------
+
+
+if __name__ == "__main__":
+    convert_dataset() # pylint: disable=no-value-for-parameter
+
+# ----------------------------------------------------------------------------

stylegan3-fun/discriminator_synthesis.py

@@ -0,0 +1,1007 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.autograd import Variable
+from torchvision import transforms
+
+import PIL
+from PIL import Image
+
+try:
+    import ffmpeg
+except ImportError:
+    raise ImportError('ffmpeg-python not found! Install it via "pip install ffmpeg-python"')
+
+import scipy.ndimage as nd
+import numpy as np
+
+import os
+import click
+from typing import Union, Tuple, Optional, List, Type
+from tqdm import tqdm
+import re
+
+from torch_utils import gen_utils
+from network_features import DiscriminatorFeatures
+
+
+# ----------------------------------------------------------------------------
+
+
+@click.group()
+def main():
+    pass
+
+
+# ----------------------------------------------------------------------------
+
+
+def get_available_layers(max_resolution: int) -> List[str]:
+    """Helper function to get the available layers given a max resolution (first block in the Discriminator)"""
+    max_res_log2 = int(np.log2(max_resolution))
+    block_resolutions = [2**i for i in range(max_res_log2, 2, -1)]
+
+    available_layers = ['from_rgb']
+    for block_res in block_resolutions:
+        # We don't add the skip layer, as it's the same as conv1 (due to in-place addition; could be changed)
+        available_layers.extend([f'b{block_res}_conv0', f'b{block_res}_conv1'])
+    # We also skip 'b4_mbstd', as it doesn't add any new information compared to b8_conv1
+    available_layers.extend(['b4_conv', 'fc', 'out'])
+    return available_layers
+
+
+# ----------------------------------------------------------------------------
+# DeepDream code; modified from Erik Linder-Norén's repository: https://github.com/eriklindernoren/PyTorch-Deep-Dream
+
+def get_image(seed: int = 0,
+              image_noise: str = 'random',
+              starting_image: Union[str, os.PathLike] = None,
+              image_size: int = 1024,
+              convert_to_grayscale: bool = False,
+              device: torch.device = torch.device('cpu')) -> Tuple[PIL.Image.Image, str]:
+    """Set the random seed (NumPy + PyTorch), as well as get an image from a path or generate a random one with the seed"""
+    torch.manual_seed(seed)
+    rnd = np.random.RandomState(seed)
+
+    # Load image or generate a random one if none is provided
+    if starting_image is not None:
+        image = Image.open(starting_image).convert('RGB').resize((image_size, image_size), Image.LANCZOS)
+    else:
+        if image_noise == 'random':
+            starting_image = f'random_image-seed_{seed:08d}.jpg'
+            image = Image.fromarray(rnd.randint(0, 255, (image_size, image_size, 3), dtype='uint8'))
+        elif image_noise == 'perlin':
+            try:
+                # Graciously using Mathieu Duchesneau's implementation: https://github.com/duchesneaumathieu/pyperlin
+                from pyperlin import FractalPerlin2D
+                starting_image = f'perlin_image-seed_{seed:08d}.jpg'
+                shape = (3, image_size, image_size)
+                resolutions = [(2**i, 2**i) for i in range(1, 6+1)]  # for lacunarity = 2.0  # TODO: set as cli variable
+                factors = [0.5**i for i in range(6)]  # for persistence = 0.5 TODO: set as cli variables
+                g_cuda = torch.Generator(device=device).manual_seed(seed)
+                rgb = FractalPerlin2D(shape, resolutions, factors, generator=g_cuda)().cpu().numpy()
+                rgb = (255 * (rgb + 1) / 2).astype(np.uint8)  # [-1.0, 1.0] => [0, 255]
+                image = Image.fromarray(rgb.transpose(1, 2, 0), 'RGB')  # Reshape leads us to weird tiling
+
+            except ImportError:
+                raise ImportError('pyperlin not found! Install it via "pip install pyperlin"')
+
+    if convert_to_grayscale:
+        image = image.convert('L').convert('RGB')  # We do a little trolling to Pillow (so we have a 3-channel image)
+
+    return image, starting_image
+
+
+def crop_resize_rotate(img: PIL.Image.Image,
+                       crop_size: int = None,
+                       new_size: int = None,
+                       rotation_deg: float = None,
+                       translate_x: float = 0.0,
+                       translate_y: float = 0.0) -> PIL.Image.Image:
+    """Center-crop the input image into a square of sides crop_size; can be resized to new_size; rotated rotation_deg counter-clockwise"""
+    # Center-crop the input image
+    if crop_size is not None:
+        w, h = img.size                                         # Input image width and height
+        img = img.crop(box=((w - crop_size) // 2,               # Left pixel coordinate
+                            (h - crop_size) // 2,               # Upper pixel coordinate
+                            (w + crop_size) // 2,               # Right pixel coordinate
+                            (h + crop_size) // 2))              # Lower pixel coordinate
+    # Resize
+    if new_size is not None:
+        img = img.resize(size=(new_size, new_size),             # Requested size of the image in pixels; (width, height)
+                         resample=Image.LANCZOS)                # Resampling filter
+    # Rotation and translation
+    if rotation_deg is not None:
+        img = img.rotate(angle=rotation_deg,                    # Angle to rotate image, counter-clockwise
+                         resample=Image.BICUBIC,                # Resampling filter; options: Image.Resampling.{NEAREST, BILINEAR, BICUBIC}
+                         expand=False,                          # If True, the whole rotated image will be shown
+                         translate=(translate_x, translate_y),  # Translate the image, from top-left corner (post-rotation)
+                         fillcolor=(0, 0, 0))                   # Black background
+    # TODO: tile the background
+    return img
+
+
+mean = np.array([0.485, 0.456, 0.406])
+std = np.array([0.229, 0.224, 0.225])
+
+preprocess = transforms.Compose([transforms.ToTensor(), transforms.Normalize(mean, std)])
+
+
+def deprocess(image_np: torch.Tensor) -> np.ndarray:
+    image_np = image_np.squeeze().transpose(1, 2, 0)
+    image_np = image_np * std.reshape((1, 1, 3)) + mean.reshape((1, 1, 3))
+    # image_np = (image_np + 1.0) / 2.0
+    image_np = np.clip(image_np, 0.0, 1.0)
+    image_np = (255 * image_np).astype('uint8')
+    return image_np
+
+
+def clip(image_tensor: torch.Tensor) -> torch.Tensor:
+    """Clamp per channel"""
+    for c in range(3):
+        m, s = mean[c], std[c]
+        image_tensor[0, c] = torch.clamp(image_tensor[0, c], -m / s, (1 - m) / s)
+    return image_tensor
+
+
+def dream(image: PIL.Image.Image,
+          model: torch.nn.Module,
+          layers: List[str],
+          channels: List[int] = None,
+          normed: bool = False,
+          sqrt_normed: bool = False,
+          iterations: int = 20,
+          lr: float = 1e-2) -> np.ndarray:
+    """ Updates the image to maximize outputs for n iterations """
+    Tensor = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor
+    image = Variable(Tensor(image), requires_grad=True)
+    for i in range(iterations):
+        model.zero_grad()
+        out = model.get_layers_features(image, layers=layers, channels=channels, normed=normed, sqrt_normed=sqrt_normed)
+        loss = sum(layer.norm() for layer in out)                   # More than one layer may be used
+        loss.backward()
+        avg_grad = np.abs(image.grad.data.cpu().numpy()).mean()
+        norm_lr = lr / avg_grad
+        image.data += norm_lr * image.grad.data
+        image.data = clip(image.data)
+        # image.data = torch.clamp(image.data, -1.0, 1.0)
+        image.grad.data.zero_()
+    return image.cpu().data.numpy()
+
+
+def deep_dream(image: PIL.Image.Image,
+               model: torch.nn.Module,
+               model_resolution: int,
+               layers: List[str],
+               channels: List[int],
+               seed: Union[int, Type[None]],
+               normed: bool,
+               sqrt_normed: bool,
+               iterations: int,
+               lr: float,
+               octave_scale: float,
+               num_octaves: int,
+               unzoom_octave: bool = False,
+               disable_inner_tqdm: bool = False,
+               ignore_initial_transform: bool = False) -> np.ndarray:
+    """ Main deep dream method """
+    # Center-crop and resize
+    if not ignore_initial_transform:
+        image = crop_resize_rotate(img=image, crop_size=min(image.size), new_size=model_resolution)
+    # Preprocess image
+    image = preprocess(image)
+    # image = torch.from_numpy(np.array(image)).permute(-1, 0, 1) / 127.5 - 1.0  # alternative
+    image = image.unsqueeze(0).cpu().data.numpy()
+    # Extract image representations for each octave
+    octaves = [image]
+    for _ in range(num_octaves - 1):
+        # Alternatively, see if we get better results with: https://www.tensorflow.org/tutorials/generative/deepdream#taking_it_up_an_octave
+        octave = nd.zoom(octaves[-1], (1, 1, 1 / octave_scale, 1 / octave_scale), order=1)
+        # Necessary for StyleGAN's Discriminator, as it cannot handle any image size
+        if unzoom_octave:
+            octave = nd.zoom(octave, np.array(octaves[-1].shape) / np.array(octave.shape), order=1)
+        octaves.append(octave)
+
+    detail = np.zeros_like(octaves[-1])
+    tqdm_desc = f'Dreaming w/layers {"|".join(x for x in layers)}'
+    tqdm_desc = f'Seed: {seed} - {tqdm_desc}' if seed is not None else tqdm_desc
+    for octave, octave_base in enumerate(tqdm(octaves[::-1], desc=tqdm_desc, disable=disable_inner_tqdm)):
+        if octave > 0:
+            # Upsample detail to new octave dimension
+            detail = nd.zoom(detail, np.array(octave_base.shape) / np.array(detail.shape), order=1)
+        # Add deep dream detail from previous octave to new base
+        input_image = octave_base + detail
+        # Get new deep dream image
+        dreamed_image = dream(input_image, model, layers, channels, normed, sqrt_normed, iterations, lr)
+        # Extract deep dream details
+        detail = dreamed_image - octave_base
+
+    return deprocess(dreamed_image)
+
+
+# ----------------------------------------------------------------------------
+
+# Helper functions (all base code taken from: https://pytorch.org/tutorials/advanced/neural_style_tutorial.html)
+
+
+class ContentLoss(nn.Module):
+
+    def __init__(self, target,):
+        super(ContentLoss, self).__init__()
+        # we 'detach' the target content from the tree used
+        # to dynamically compute the gradient: this is a stated value,
+        # not a variable. Otherwise the forward method of the criterion
+        # will throw an error.
+        self.target = target.detach()
+
+    def forward(self, input):
+        self.loss = F.mse_loss(input, self.target)
+        return input
+
+
+def gram_matrix(input):
+    a, b, c, d = input.size()  # (batch_size, no. feature maps, dims of a f. map (N=c*d))
+
+    features = input.view(a * b, c * d)  # resize F_XL into \hat F_XL
+
+    G = torch.mm(features, features.t())  # compute the gram product
+
+    # 'Normalize' the values of the gram matrix by dividing by the number of element in each feature maps.
+    return G.div(a * b * c * d)  # can also do torch.numel(input) to get the number of elements
+
+
+class StyleLoss(nn.Module):
+    def __init__(self, target_feature):
+        super(StyleLoss, self).__init__()
+        self.target = gram_matrix(target_feature).detach()
+
+    def forward(self, input):
+        G = gram_matrix(input)
+        self.loss = F.mse_loss(G, self.target)
+        return input
+
+
+@main.command(name='style-transfer', help='Use the StyleGAN2/3 Discriminator to perform style transfer')
+@click.option('--network', 'network_pkl', help='Network pickle filename', required=True)
+@click.option('--cfg', type=click.Choice(['stylegan3-t', 'stylegan3-r', 'stylegan2']), help='Model base configuration', default=None)
+@click.option('--content', type=str, help='Content image filename (url or local path)', required=True)
+@click.option('--style', type=str, help='Style image filename (url or local path)', required=True)
+def style_transfer_discriminator(
+        ctx: click.Context,
+        network_pkl: str,
+        cfg: str,
+        content: str,
+        style: str,
+):
+    print('Coming soon!')
+    # Reference: https://pytorch.org/tutorials/advanced/neural_style_tutorial.html
+
+    # Set up device
+    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+
+    imsize = 512 if torch.cuda.is_available() else 128  # use small size if no gpu
+
+    loader = transforms.Compose([transforms.Resize(imsize),  # scale imported image
+                                transforms.ToTensor()])  # transform it into a torch tensor
+
+    # Helper function
+    def image_loader(image_name):
+        image = Image.open(image_name)
+        # fake batch dimension required to fit network's input dimensions
+        image = loader(image).unsqueeze(0)
+        return image.to(device, torch.float)
+
+    style_img = image_loader(style)
+    content_img = image_loader(content)
+
+    # This shouldn't really happen, but just in case
+    assert style_img.size() == content_img.size(), 'Style and content images must be the same size'
+
+    unloader = transforms.ToPILImage()  # reconvert into PIL image
+
+    # Load Discriminator
+    D = gen_utils.load_network('D', network_pkl, cfg, device)
+    # TODO: finish this!
+
+
+# ----------------------------------------------------------------------------
+
+
+@main.command(name='dream', help='Discriminator Dreaming with the StyleGAN2/3 Discriminator and the chosen layers')
+@click.pass_context
+@click.option('--network', 'network_pkl', help='Network pickle filename', required=True)
+@click.option('--cfg', type=click.Choice(['stylegan3-t', 'stylegan3-r', 'stylegan2']), help='Model base configuration', default=None)
+# Synthesis options
+@click.option('--seeds', type=gen_utils.num_range, help='Random seeds to use. Accepted comma-separated values, ranges, or combinations: "a,b,c", "a-c", "a,b-d,e".', default='0')
+@click.option('--random-image-noise', '-noise', 'image_noise', type=click.Choice(['random', 'perlin']), default='perlin', show_default=True)
+@click.option('--starting-image', type=str, help='Path to image to start from', default=None)
+@click.option('--convert-to-grayscale', '-grayscale', is_flag=True, help='Add flag to grayscale the initial image')
+@click.option('--class', 'class_idx', type=int, help='Class label (unconditional if not specified)', default=None)
+@click.option('--lr', 'learning_rate', type=float, help='Learning rate', default=1e-2, show_default=True)
+@click.option('--iterations', '-it', type=int, help='Number of gradient ascent steps per octave', default=20, show_default=True)
+# Layer options
+@click.option('--layers', type=str, help='Layers of the Discriminator to use as the features. If "all", will generate a dream image per available layer in the loaded model. If "use_all", will use all available layers.', default='b16_conv1', show_default=True)
+@click.option('--channels', type=gen_utils.num_range, help='Comma-separated list and/or range of the channels of the Discriminator to use as the features. If "None", will use all channels in each specified layer.', default=None, show_default=True)
+@click.option('--normed', 'norm_model_layers', is_flag=True, help='Add flag to divide the features of each layer of D by its number of elements')
+@click.option('--sqrt-normed', 'sqrt_norm_model_layers', is_flag=True, help='Add flag to divide the features of each layer of D by the square root of its number of elements')
+# Octaves options
+@click.option('--num-octaves', type=int, help='Number of octaves', default=5, show_default=True)
+@click.option('--octave-scale', type=float, help='Image scale between octaves', default=1.4, show_default=True)
+@click.option('--unzoom-octave', type=bool, help='Set to True for the octaves to be unzoomed (this will be slower)', default=True, show_default=True)
+# Extra parameters for saving the results
+@click.option('--outdir', type=click.Path(file_okay=False), help='Directory path to save the results', default=os.path.join(os.getcwd(), 'out', 'discriminator_synthesis'), show_default=True, metavar='DIR')
+@click.option('--description', '-desc', type=str, help='Additional description name for the directory path to save results', default='', show_default=True)
+def discriminator_dream(
+        ctx: click.Context,
+        network_pkl: Union[str, os.PathLike],
+        cfg: Optional[str],
+        seeds: List[int],
+        image_noise: str,
+        starting_image: Union[str, os.PathLike],
+        convert_to_grayscale: bool,
+        class_idx: Optional[int],  # TODO: conditional model
+        learning_rate: float,
+        iterations: int,
+        layers: str,
+        channels: Optional[List[int]],
+        norm_model_layers: bool,
+        sqrt_norm_model_layers: bool,
+        num_octaves: int,
+        octave_scale: float,
+        unzoom_octave: bool,
+        outdir: Union[str, os.PathLike],
+        description: str,
+):
+    # Set up device
+    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+
+    # Load Discriminator
+    D = gen_utils.load_network('D', network_pkl, cfg, device)
+
+    # Get the model resolution (image resizing and getting available layers)
+    model_resolution = D.img_resolution
+
+    # TODO: do this better, as we can combine these conditions later
+    layers = layers.split(',')
+
+    # We will use the features of the Discriminator, on the layer specified by the user
+    model = DiscriminatorFeatures(D).requires_grad_(False).to(device)
+
+    if 'all' in layers:
+        # Get all the available layers in a list
+        layers = get_available_layers(max_resolution=model_resolution)
+
+        for seed in seeds:
+            # Get the image and image name
+            image, starting_image = get_image(seed=seed, image_noise=image_noise,
+                                              starting_image=starting_image,
+                                              image_size=model_resolution,
+                                              convert_to_grayscale=convert_to_grayscale)
+
+            # Make the run dir in the specified output directory
+            desc = f'discriminator-dream-all_layers-seed_{seed}'
+            desc = f'{desc}-{description}' if len(description) != 0 else desc
+            run_dir = gen_utils.make_run_dir(outdir, desc)
+
+            # Save starting image
+            image.save(os.path.join(run_dir, f'{os.path.basename(starting_image).split(".")[0]}.jpg'))
+
+            # Save the configuration used
+            ctx.obj = {
+                'network_pkl': network_pkl,
+                'synthesis_options': {
+                    'seed': seed,
+                    'random_image_noise': image_noise,
+                    'starting_image': starting_image,
+                    'class_idx': class_idx,
+                    'learning_rate': learning_rate,
+                    'iterations': iterations},
+                'layer_options': {
+                    'layer': layers,
+                    'channels': channels,
+                    'norm_model_layers': norm_model_layers,
+                    'sqrt_norm_model_layers': sqrt_norm_model_layers},
+                'octaves_options': {
+                    'num_octaves': num_octaves,
+                    'octave_scale': octave_scale,
+                    'unzoom_octave': unzoom_octave},
+                'extra_parameters': {
+                    'outdir': run_dir,
+                    'description': description}
+            }
+            # Save the run configuration
+            gen_utils.save_config(ctx=ctx, run_dir=run_dir)
+
+            # For each layer:
+            for layer in layers:
+                # Extract deep dream image
+                dreamed_image = deep_dream(image, model, model_resolution, layers=[layer], channels=channels, seed=seed, normed=norm_model_layers,
+                                           sqrt_normed=sqrt_norm_model_layers, iterations=iterations, lr=learning_rate,
+                                           octave_scale=octave_scale, num_octaves=num_octaves, unzoom_octave=unzoom_octave)
+
+                # Save the resulting dreamed image
+                filename = f'layer-{layer}_dreamed_{os.path.basename(starting_image).split(".")[0]}.jpg'
+                Image.fromarray(dreamed_image, 'RGB').save(os.path.join(run_dir, filename))
+
+    else:
+        if 'use_all' in layers:
+            # Get all available layers
+            layers = get_available_layers(max_resolution=model_resolution)
+        else:
+            # Parse the layers given by the user and leave only those available by the model
+            available_layers = get_available_layers(max_resolution=model_resolution)
+            layers = [layer for layer in layers if layer in available_layers]
+
+        # Make the run dir in the specified output directory
+        desc = f'discriminator-dream-layers_{"-".join(x for x in layers)}'
+        desc = f'{desc}-{description}' if len(description) != 0 else desc
+        run_dir = gen_utils.make_run_dir(outdir, desc)
+
+        starting_images, used_seeds = [], []
+        for seed in seeds:
+            # Get the image and image name
+            image, starting_image = get_image(seed=seed, image_noise=image_noise,
+                                              starting_image=starting_image,
+                                              image_size=model_resolution,
+                                              convert_to_grayscale=convert_to_grayscale)
+
+            # Extract deep dream image
+            dreamed_image = deep_dream(image, model, model_resolution, layers=layers, channels=channels, seed=seed, normed=norm_model_layers,
+                                       sqrt_normed=sqrt_norm_model_layers, iterations=iterations, lr=learning_rate,
+                                       octave_scale=octave_scale, num_octaves=num_octaves, unzoom_octave=unzoom_octave)
+
+            # For logging later
+            starting_images.append(starting_image)
+            used_seeds.append(seed)
+
+            # Save the resulting image and initial image
+            filename = f'dreamed_{os.path.basename(starting_image)}'
+            Image.fromarray(dreamed_image, 'RGB').save(os.path.join(run_dir, filename))
+            image.save(os.path.join(run_dir, os.path.basename(starting_image)))
+            starting_image = None
+
+        # Save the configuration used
+        ctx.obj = {
+            'network_pkl': network_pkl,
+            'synthesis_options': {
+                'seeds': used_seeds,
+                'starting_image': starting_images,
+                'class_idx': class_idx,
+                'learning_rate': learning_rate,
+                'iterations': iterations},
+            'layer_options': {
+                'layer': layers,
+                'channels': channels,
+                'norm_model_layers': norm_model_layers,
+                'sqrt_norm_model_layers': sqrt_norm_model_layers},
+            'octaves_options': {
+                'octave_scale': octave_scale,
+                'num_octaves': num_octaves,
+                'unzoom_octave': unzoom_octave},
+            'extra_parameters': {
+                'outdir': run_dir,
+                'description': description}
+        }
+        # Save the run configuration
+        gen_utils.save_config(ctx=ctx, run_dir=run_dir)
+
+
+# ----------------------------------------------------------------------------
+
+
+@main.command(name='dream-zoom',
+              help='Zoom/rotate/translate after each Discriminator Dreaming iteration. A video will be saved.')
+@click.pass_context
+@click.option('--network', 'network_pkl', help='Network pickle filename', required=True)
+@click.option('--cfg', type=click.Choice(['stylegan3-t', 'stylegan3-r', 'stylegan2']), help='Model base configuration', default=None)
+# Synthesis options
+@click.option('--seed', type=int, help='Random seed to use', default=0, show_default=True)
+@click.option('--random-image-noise', '-noise', 'image_noise', type=click.Choice(['random', 'perlin']), default='random', show_default=True)
+@click.option('--starting-image', type=str, help='Path to image to start from', default=None)
+@click.option('--convert-to-grayscale', '-grayscale', is_flag=True, help='Add flag to grayscale the initial image')
+@click.option('--class', 'class_idx', type=int, help='Class label (unconditional if not specified)', default=None)
+@click.option('--lr', 'learning_rate', type=float, help='Learning rate', default=5e-3, show_default=True)
+@click.option('--iterations', '-it', type=click.IntRange(min=1), help='Number of gradient ascent steps per octave', default=10, show_default=True)
+# Layer options
+@click.option('--layers', type=str, help='Comma-separated list of the layers of the Discriminator to use as the features. If "use_all", will use all available layers.', default='b16_conv0', show_default=True)
+@click.option('--channels', type=gen_utils.num_range, help='Comma-separated list and/or range of the channels of the Discriminator to use as the features. If "None", will use all channels in each specified layer.', default=None, show_default=True)
+@click.option('--normed', 'norm_model_layers', is_flag=True, help='Add flag to divide the features of each layer of D by its number of elements')
+@click.option('--sqrt-normed', 'sqrt_norm_model_layers', is_flag=True, help='Add flag to divide the features of each layer of D by the square root of its number of elements')
+# Octaves options
+@click.option('--num-octaves', type=click.IntRange(min=1), help='Number of octaves', default=5, show_default=True)
+@click.option('--octave-scale', type=float, help='Image scale between octaves', default=1.4, show_default=True)
+@click.option('--unzoom-octave', type=bool, help='Set to True for the octaves to be unzoomed (this will be slower)', default=False, show_default=True)
+# Individual frame manipulation options
+@click.option('--pixel-zoom', '-zoom', type=int, help='How many pixels to zoom per step (positive for zoom in, negative for zoom out, padded with black)', default=2, show_default=True)
+@click.option('--rotation-deg', '-rot', type=float, help='Rotate image counter-clockwise per frame (padded with black)', default=0.0, show_default=True)
+@click.option('--translate-x', '-tx', type=float, help='Translate the image in the horizontal axis per frame (from left to right, padded with black)', default=0.0, show_default=True)
+@click.option('--translate-y', '-ty', type=float, help='Translate the image in the vertical axis per frame (from top to bottom, padded with black)', default=0.0, show_default=True)
+# Video options
+@click.option('--fps', type=gen_utils.parse_fps, help='FPS for the mp4 video of optimization progress (if saved)', default=25, show_default=True)
+@click.option('--duration-sec', type=float, help='Duration length of the video', default=15.0, show_default=True)
+@click.option('--reverse-video', is_flag=True, help='Add flag to reverse the generated video')
+@click.option('--include-starting-image', type=bool, help='Include the starting image in the final video', default=True, show_default=True)
+# Extra parameters for saving the results
+@click.option('--outdir', type=click.Path(file_okay=False), help='Directory path to save the results', default=os.path.join(os.getcwd(), 'out', 'discriminator_synthesis'), show_default=True, metavar='DIR')
+@click.option('--description', '-desc', type=str, help='Additional description name for the directory path to save results', default='', show_default=True)
+def discriminator_dream_zoom(
+        ctx: click.Context,
+        network_pkl: Union[str, os.PathLike],
+        cfg: Optional[str],
+        seed: int,
+        image_noise: Optional[str],
+        starting_image: Optional[Union[str, os.PathLike]],
+        convert_to_grayscale: bool,
+        class_idx: Optional[int],  # TODO: conditional model
+        learning_rate: float,
+        iterations: int,
+        layers: str,
+        channels: List[int],
+        norm_model_layers: Optional[bool],
+        sqrt_norm_model_layers: Optional[bool],
+        num_octaves: int,
+        octave_scale: float,
+        unzoom_octave: Optional[bool],
+        pixel_zoom: int,
+        rotation_deg: float,
+        translate_x: int,
+        translate_y: int,
+        fps: int,
+        duration_sec: float,
+        reverse_video: bool,
+        include_starting_image: bool,
+        outdir: Union[str, os.PathLike],
+        description: str,
+):
+    # Set up device
+    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+
+    # Load Discriminator
+    D = gen_utils.load_network('D', network_pkl, cfg, device)
+
+    # Get the model resolution (for resizing the starting image if needed)
+    model_resolution = D.img_resolution
+    zoom_size = model_resolution - 2 * pixel_zoom
+
+    layers = layers.split(',')
+    if 'use_all' in layers:
+        # Get all available layers
+        layers = get_available_layers(max_resolution=model_resolution)
+    else:
+        # Parse the layers given by the user and leave only those available by the model
+        available_layers = get_available_layers(max_resolution=model_resolution)
+        layers = [layer for layer in layers if layer in available_layers]
+
+    # We will use the features of the Discriminator, on the layer specified by the user
+    model = DiscriminatorFeatures(D).requires_grad_(False).to(device)
+
+    # Get the image and image name
+    image, starting_image = get_image(seed=seed, image_noise=image_noise,
+                                      starting_image=starting_image,
+                                      image_size=model_resolution,
+                                      convert_to_grayscale=convert_to_grayscale)
+
+    # Make the run dir in the specified output directory
+    desc = 'discriminator-dream-zoom'
+    desc = f'{desc}-{description}' if len(description) != 0 else desc
+    run_dir = gen_utils.make_run_dir(outdir, desc)
+
+    # Save the configuration used
+    ctx.obj = {
+        'network_pkl': network_pkl,
+        'synthesis_options': {
+            'seed': seed,
+            'random_image_noise': image_noise,
+            'starting_image': starting_image,
+            'class_idx': class_idx,
+            'learning_rate': learning_rate,
+            'iterations': iterations
+        },
+        'layer_options': {
+            'layers': layers,
+            'channels': channels,
+            'norm_model_layers': norm_model_layers,
+            'sqrt_norm_model_layers': sqrt_norm_model_layers
+        },
+        'octaves_options': {
+            'num_octaves': num_octaves,
+            'octave_scale': octave_scale,
+            'unzoom_octave': unzoom_octave
+        },
+        'frame_manipulation_options': {
+            'pixel_zoom': pixel_zoom,
+            'rotation_deg': rotation_deg,
+            'translate_x': translate_x,
+            'translate_y': translate_y,
+        },
+        'video_options': {
+            'fps': fps,
+            'duration_sec': duration_sec,
+            'reverse_video': reverse_video,
+            'include_starting_image': include_starting_image,
+        },
+        'extra_parameters': {
+            'outdir': run_dir,
+            'description': description
+        }
+    }
+    # Save the run configuration
+    gen_utils.save_config(ctx=ctx, run_dir=run_dir)
+
+    num_frames = int(np.rint(duration_sec * fps))  # Number of frames for the video
+    n_digits = int(np.log10(num_frames)) + 1       # Number of digits for naming each frame
+
+    # Save the starting image
+    starting_image_name = f'dreamed_{0:0{n_digits}d}.jpg' if include_starting_image else 'starting_image.jpg'
+    image.save(os.path.join(run_dir, starting_image_name))
+
+    for idx, frame in enumerate(tqdm(range(num_frames), desc='Dreaming...', unit='frame')):
+        # Zoom in after the first frame
+        if idx > 0:
+            image = crop_resize_rotate(image, crop_size=zoom_size, new_size=model_resolution,
+                                       rotation_deg=rotation_deg, translate_x=translate_x, translate_y=translate_y)
+        # Extract deep dream image
+        dreamed_image = deep_dream(image, model, model_resolution, layers=layers, seed=seed, normed=norm_model_layers,
+                                   sqrt_normed=sqrt_norm_model_layers, iterations=iterations, channels=channels,
+                                   lr=learning_rate, octave_scale=octave_scale, num_octaves=num_octaves,
+                                   unzoom_octave=unzoom_octave, disable_inner_tqdm=True)
+
+        # Save the resulting image and initial image
+        filename = f'dreamed_{idx + 1:0{n_digits}d}.jpg'
+        Image.fromarray(dreamed_image, 'RGB').save(os.path.join(run_dir, filename))
+
+        # Now, the dreamed image is the starting image
+        image = Image.fromarray(dreamed_image, 'RGB')
+
+    # Save the final video
+    gen_utils.save_video_from_images(run_dir=run_dir, image_names=f'dreamed_%0{n_digits}d.jpg',
+                                     video_name='dream-zoom', fps=fps, reverse_video=reverse_video)
+
+
+# ----------------------------------------------------------------------------
+
+@main.command(name='channel-zoom', help='Dream zoom using only the specified channels in the selected layer')
+@click.pass_context
+@click.option('--network', 'network_pkl', help='Network pickle filename', required=True)
+@click.option('--cfg', type=click.Choice(['stylegan3-t', 'stylegan3-r', 'stylegan2']), help='Model base configuration', default=None)
+# Synthesis options
+@click.option('--seed', type=int, help='Random seed to use', default=0, show_default=True)
+@click.option('--random-image-noise', '-noise', 'image_noise', type=click.Choice(['random', 'perlin']), default='random', show_default=True)
+@click.option('--starting-image', type=str, help='Path to image to start from', default=None)
+@click.option('--convert-to-grayscale', '-grayscale', is_flag=True, help='Add flag to grayscale the initial image')
+@click.option('--class', 'class_idx', type=int, help='Class label (unconditional if not specified)', default=None)
+@click.option('--lr', 'learning_rate', type=float, help='Learning rate', default=5e-3, show_default=True)
+@click.option('--iterations', '-it', type=click.IntRange(min=1), help='Number of gradient ascent steps per octave', default=10, show_default=True)
+# Layer options
+@click.option('--layer', type=str, help='Layers of the Discriminator to use as the features.', default='b8_conv0', show_default=True)
+@click.option('--normed', 'norm_model_layers', is_flag=True, help='Add flag to divide the features of each layer of D by its number of elements')
+@click.option('--sqrt-normed', 'sqrt_norm_model_layers', is_flag=True, help='Add flag to divide the features of each layer of D by the square root of its number of elements')
+# Octaves options
+@click.option('--num-octaves', type=click.IntRange(min=1), help='Number of octaves', default=5, show_default=True)
+@click.option('--octave-scale', type=float, help='Image scale between octaves', default=1.4, show_default=True)
+@click.option('--unzoom-octave', type=bool, help='Set to True for the octaves to be unzoomed (this will be slower)', default=False, show_default=True)
+# Individual frame manipulation options
+@click.option('--pixel-zoom', '-zoom', type=int, help='How many pixels to zoom per step (positive for zoom in, negative for zoom out, padded with black)', default=2, show_default=True)
+@click.option('--rotation-deg', '-rot', type=float, help='Rotate image counter-clockwise per frame (padded with black)', default=0.0, show_default=True)
+@click.option('--translate-x', '-tx', type=float, help='Translate the image in the horizontal axis per frame (from left to right, padded with black)', default=0.0, show_default=True)
+@click.option('--translate-y', '-ty', type=float, help='Translate the image in the vertical axis per frame (from top to bottom, padded with black)', default=0.0, show_default=True)
+# Video options
+@click.option('--frames-per-channel', type=click.IntRange(min=1), help='Number of frames per channel', default=1, show_default=True)
+@click.option('--fps', type=gen_utils.parse_fps, help='FPS for the mp4 video of optimization progress (if saved)', default=25, show_default=True)
+@click.option('--reverse-video', is_flag=True, help='Add flag to reverse the generated video')
+@click.option('--include-starting-image', type=bool, help='Include the starting image in the final video', default=True, show_default=True)
+# Extra parameters for saving the results
+@click.option('--outdir', type=click.Path(file_okay=False), help='Directory path to save the results', default=os.path.join(os.getcwd(), 'out', 'discriminator_synthesis'), show_default=True, metavar='DIR')
+@click.option('--description', '-desc', type=str, help='Additional description name for the directory path to save results', default='', show_default=True)
+def channel_zoom(
+        ctx: click.Context,
+        network_pkl: Union[str, os.PathLike],
+        cfg: Optional[str],
+        seed: int,
+        image_noise: Optional[str],
+        starting_image: Optional[Union[str, os.PathLike]],
+        convert_to_grayscale: bool,
+        class_idx: Optional[int],  # TODO: conditional model
+        learning_rate: float,
+        iterations: int,
+        layer: str,
+        norm_model_layers: Optional[bool],
+        sqrt_norm_model_layers: Optional[bool],
+        num_octaves: int,
+        octave_scale: float,
+        unzoom_octave: Optional[bool],
+        pixel_zoom: int,
+        rotation_deg: float,
+        translate_x: int,
+        translate_y: int,
+        frames_per_channel: int,
+        fps: int,
+        reverse_video: bool,
+        include_starting_image: bool,
+        outdir: Union[str, os.PathLike],
+        description: str,
+):
+    """Zoom in using all the channels of a network (or a specified layer)"""
+    # Set up device
+    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+
+    # Load Discriminator
+    D = gen_utils.load_network('D', network_pkl, cfg, device)
+
+    # Get the model resolution (for resizing the starting image if needed)
+    model_resolution = D.img_resolution
+    zoom_size = model_resolution - 2 * pixel_zoom
+
+    if 'use_all' in layer:
+        ctx.fail('Cannot use "use_all" with this command. Please specify the layers you want to use.')
+    else:
+        # Parse the layers given by the user and leave only those available by the model
+        available_layers = get_available_layers(max_resolution=model_resolution)
+        assert layer in available_layers, f'Layer {layer} not available. Available layers: {available_layers}'
+        layers = [layer]
+
+    # We will use the features of the Discriminator, on the layer specified by the user
+    model = DiscriminatorFeatures(D).requires_grad_(False).to(device)
+
+    # Get the image and image name
+    image, starting_image = get_image(seed=seed, image_noise=image_noise,
+                                      starting_image=starting_image,
+                                      image_size=model_resolution,
+                                      convert_to_grayscale=convert_to_grayscale)
+
+    # Make the run dir in the specified output directory
+    desc = 'discriminator-channel-zoom'
+    desc = f'{desc}-{description}' if len(description) != 0 else desc
+    run_dir = gen_utils.make_run_dir(outdir, desc)
+
+    # Finally, let's get the number of channels in the selected layer
+    channels_dict = {res: D.get_submodule(f'b{res}.conv0').out_channels for res in D.block_resolutions}
+    channels_dict[4] = D.get_submodule('b4.conv').out_channels  # Last block has a different name
+    # Get the dimension of the block from the selected layer (e.g., from 'b128_conv0' get '128')
+    block_resolution = re.search(r'b(\d+)_', layer).group(1)
+    total_channels = channels_dict[int(block_resolution)]
+    # Make a list of all the channels, each repeated frames_per_channel
+    channels = np.repeat(np.arange(total_channels), frames_per_channel)
+
+    num_frames = int(np.rint(total_channels * frames_per_channel))  # Number of frames for the video
+    n_digits = int(np.log10(num_frames)) + 1  # Number of digits for naming each frame
+
+    # Save the starting image
+    starting_image_name = f'dreamed_{0:0{n_digits}d}.jpg' if include_starting_image else 'starting_image.jpg'
+    image.save(os.path.join(run_dir, starting_image_name))
+
+    for idx, frame in enumerate(tqdm(range(num_frames), desc='Dreaming...', unit='frame')):
+        # Zoom in after the first frame
+        if idx > 0:
+            image = crop_resize_rotate(image, crop_size=zoom_size, new_size=model_resolution,
+                                       rotation_deg=rotation_deg, translate_x=translate_x, translate_y=translate_y)
+        # Extract deep dream image
+        dreamed_image = deep_dream(image, model, model_resolution, layers=layers, seed=seed, normed=norm_model_layers,
+                                   sqrt_normed=sqrt_norm_model_layers, iterations=iterations, channels=channels[idx:idx + 1],
+                                   lr=learning_rate, octave_scale=octave_scale, num_octaves=num_octaves,
+                                   unzoom_octave=unzoom_octave, disable_inner_tqdm=True)
+
+        # Save the resulting image and initial image
+        filename = f'dreamed_{idx + 1:0{n_digits}d}.jpg'
+        Image.fromarray(dreamed_image, 'RGB').save(os.path.join(run_dir, filename))
+
+        # Now, the dreamed image is the starting image
+        image = Image.fromarray(dreamed_image, 'RGB')
+
+    # Save the final video
+    gen_utils.save_video_from_images(run_dir=run_dir, image_names=f'dreamed_%0{n_digits}d.jpg', video_name='channel-zoom',
+                                     fps=fps, reverse_video=reverse_video)
+
+    # Save the configuration used
+    ctx.obj = {
+        'network_pkl': network_pkl,
+        'synthesis_options': {
+            'seed': seed,
+            'random_image_noise': image_noise,
+            'starting_image': starting_image,
+            'class_idx': class_idx,
+            'learning_rate': learning_rate,
+            'iterations': iterations
+        },
+        'layer_options': {
+            'layer': layer,
+            'channels': 'all',
+            'total_channels': total_channels,
+            'norm_model_layers': norm_model_layers,
+            'sqrt_norm_model_layers': sqrt_norm_model_layers
+        },
+        'octaves_options': {
+            'num_octaves': num_octaves,
+            'octave_scale': octave_scale,
+            'unzoom_octave': unzoom_octave
+        },
+        'frame_manipulation_options': {
+            'pixel_zoom': pixel_zoom,
+            'rotation_deg': rotation_deg,
+            'translate_x': translate_x,
+            'translate_y': translate_y,
+        },
+        'video_options': {
+            'fps': fps,
+            'frames_per_channel': frames_per_channel,
+            'reverse_video': reverse_video,
+            'include_starting_image': include_starting_image,
+        },
+        'extra_parameters': {
+            'outdir': run_dir,
+            'description': description
+        }
+    }
+    # Save the run configuration
+    gen_utils.save_config(ctx=ctx, run_dir=run_dir)
+
+
+# ----------------------------------------------------------------------------
+
+
+@main.command(name='interp', help='Interpolate between two or more seeds')
+@click.pass_context
+@click.option('--network', 'network_pkl', help='Network pickle filename', required=True)
+@click.option('--cfg', type=click.Choice(['stylegan3-t', 'stylegan3-r', 'stylegan2']), help='Model base configuration', default=None)
+# Synthesis options
+@click.option('--seeds', type=gen_utils.num_range, help='Random seeds to generate the Perlin noise from', required=True)
+@click.option('--interp-type', '-interp', type=click.Choice(['linear', 'spherical']), help='Type of interpolation in Z or W', default='spherical', show_default=True)
+@click.option('--smooth', is_flag=True, help='Add flag to smooth the interpolation between the seeds')
+@click.option('--random-image-noise', '-noise', 'image_noise', type=click.Choice(['random', 'perlin']), default='random', show_default=True)
+@click.option('--starting-image', type=str, help='Path to image to start from', default=None)
+@click.option('--convert-to-grayscale', '-grayscale', is_flag=True, help='Add flag to grayscale the initial image')
+@click.option('--class', 'class_idx', type=int, help='Class label (unconditional if not specified)', default=None)
+@click.option('--lr', 'learning_rate', type=float, help='Learning rate', default=5e-3, show_default=True)
+@click.option('--iterations', '-it', type=click.IntRange(min=1), help='Number of gradient ascent steps per octave', default=10, show_default=True)
+# Layer options
+@click.option('--layers', type=str, help='Comma-separated list of the layers of the Discriminator to use as the features. If "use_all", will use all available layers.', default='b16_conv0', show_default=True)
+@click.option('--channels', type=gen_utils.num_range, help='Comma-separated list and/or range of the channels of the Discriminator to use as the features. If "None", will use all channels in each specified layer.', default=None, show_default=True)
+@click.option('--normed', 'norm_model_layers', is_flag=True, help='Add flag to divide the features of each layer of D by its number of elements')
+@click.option('--sqrt-normed', 'sqrt_norm_model_layers', is_flag=True, help='Add flag to divide the features of each layer of D by the square root of its number of elements')
+# Octaves options
+@click.option('--num-octaves', type=click.IntRange(min=1), help='Number of octaves', default=5, show_default=True)
+@click.option('--octave-scale', type=float, help='Image scale between octaves', default=1.4, show_default=True)
+@click.option('--unzoom-octave', type=bool, help='Set to True for the octaves to be unzoomed (this will be slower)', default=False, show_default=True)
+# TODO: Individual frame manipulation options
+# Video options
+@click.option('--seed-sec', '-sec', type=float, help='Number of seconds between each seed transition', default=5.0, show_default=True)
+@click.option('--fps', type=gen_utils.parse_fps, help='FPS for the mp4 video of optimization progress (if saved)', default=25, show_default=True)
+# Extra parameters for saving the results
+@click.option('--outdir', type=click.Path(file_okay=False), help='Directory path to save the results', default=os.path.join(os.getcwd(), 'out', 'discriminator_synthesis'), show_default=True, metavar='DIR')
+@click.option('--description', '-desc', type=str, help='Additional description name for the directory path to save results', default='', show_default=True)
+def random_interpolation(
+        ctx: click.Context,
+        network_pkl: Union[str, os.PathLike],
+        cfg: Optional[str],
+        seeds: List[int],
+        interp_type: Optional[str],
+        smooth: Optional[bool],
+        image_noise: Optional[str],
+        starting_image: Optional[Union[str, os.PathLike]],
+        convert_to_grayscale: bool,
+        class_idx: Optional[int],  # TODO: conditional model
+        learning_rate: float,
+        iterations: int,
+        layers: str,
+        channels: List[int],
+        norm_model_layers: Optional[bool],
+        sqrt_norm_model_layers: Optional[bool],
+        num_octaves: int,
+        octave_scale: float,
+        unzoom_octave: Optional[bool],
+        seed_sec: float,
+        fps: int,
+        outdir: Union[str, os.PathLike],
+        description: str,
+):
+    """Do a latent walk between random Perlin images (given the seeds) and generate a video with these frames."""
+    # TODO: To make this better and more stable, we generate Perlin noise animations, not interpolations
+    # Set up device
+    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+
+    # Load Discriminator
+    D = gen_utils.load_network('D', network_pkl, cfg, device)
+
+    # Get model resolution
+    model_resolution = D.img_resolution
+    model = DiscriminatorFeatures(D).requires_grad_(False).to(device)
+
+    layers = layers.split(',')
+    # Get all available layers
+    if 'use_all' in layers:
+        layers = get_available_layers(max_resolution=model_resolution)
+    else:
+        # Parse the layers given by the user and leave only those available by the model
+        available_layers = get_available_layers(max_resolution=model_resolution)
+        layers = [layer for layer in layers if layer in available_layers]
+
+    # Make the run dir in the specified output directory
+    desc = f'random-interp-layers_{"-".join(x for x in layers)}'
+    desc = f'{desc}-{description}' if len(description) != 0 else desc
+    run_dir = gen_utils.make_run_dir(outdir, desc)
+
+    # Number of steps to take between each random image
+    n_steps = int(np.rint(seed_sec * fps))
+    # Total number of frames
+    num_frames = int(n_steps * (len(seeds) - 1))
+    # Total video length in seconds
+    duration_sec = num_frames / fps
+
+    # Number of digits for naming purposes
+    n_digits = int(np.log10(num_frames)) + 1
+
+    # Create interpolation of noises
+    random_images = []
+    for seed in seeds:
+        # Get the starting seed and image
+        image, _ = get_image(seed=seed, image_noise=image_noise, starting_image=starting_image,
+                             image_size=model_resolution, convert_to_grayscale=convert_to_grayscale)
+        image = np.array(image) / 255.0
+        random_images.append(image)
+    random_images = np.stack(random_images)
+
+    all_images = np.empty([0] + list(random_images.shape[1:]), dtype=np.float32)
+    # Do interpolation
+    for i in range(len(random_images) - 1):
+        # Interpolate between each pair of images
+        interp = gen_utils.interpolate(random_images[i], random_images[i + 1], n_steps, interp_type, smooth)
+        # Append it to the list of all images
+        all_images = np.append(all_images, interp, axis=0)
+
+    # DeepDream expects a list of PIL.Image objects
+    pil_images = []
+    for idx in range(len(all_images)):
+        im = (255 * all_images[idx]).astype(dtype=np.uint8)
+        pil_images.append(Image.fromarray(im))
+
+    for idx, image in enumerate(tqdm(pil_images, desc='Interpolating...', unit='frame', total=num_frames)):
+        # Extract deep dream image
+        dreamed_image = deep_dream(image, model, model_resolution, layers=layers, channels=channels, seed=None,
+                                   normed=norm_model_layers, disable_inner_tqdm=True, ignore_initial_transform=True,
+                                   sqrt_normed=sqrt_norm_model_layers, iterations=iterations, lr=learning_rate,
+                                   octave_scale=octave_scale, num_octaves=num_octaves, unzoom_octave=unzoom_octave)
+
+        # Save the resulting image and initial image
+        filename = f'{image_noise}-interpolation_frame_{idx:0{n_digits}d}.jpg'
+        Image.fromarray(dreamed_image, 'RGB').save(os.path.join(run_dir, filename))
+
+    # Save the configuration used
+    ctx.obj = {
+        'network_pkl': network_pkl,
+        'synthesis_options': {
+            'seeds': seeds,
+            'starting_image': starting_image,
+            'class_idx': class_idx,
+            'learning_rate': learning_rate,
+            'iterations': iterations},
+        'layer_options': {
+            'layer': layers,
+            'channels': channels,
+            'norm_model_layers': norm_model_layers,
+            'sqrt_norm_model_layers': sqrt_norm_model_layers},
+        'octaves_options': {
+            'octave_scale': octave_scale,
+            'num_octaves': num_octaves,
+            'unzoom_octave': unzoom_octave},
+        'extra_parameters': {
+            'outdir': run_dir,
+            'description': description}
+    }
+    # Save the run configuration
+    gen_utils.save_config(ctx=ctx, run_dir=run_dir)
+
+    # Generate video
+    print('Saving video...')
+    ffmpeg_command = r'/usr/bin/ffmpeg' if os.name != 'nt' else r'C:\\Ffmpeg\\bin\\ffmpeg.exe'
+    stream = ffmpeg.input(os.path.join(run_dir, f'{image_noise}-interpolation_frame_%0{n_digits}d.jpg'), framerate=fps)
+    stream = ffmpeg.output(stream, os.path.join(run_dir, f'{image_noise}-interpolation.mp4'), crf=20, pix_fmt='yuv420p')
+    ffmpeg.run(stream, capture_stdout=True, capture_stderr=True, cmd=ffmpeg_command)
+
+# ----------------------------------------------------------------------------
+
+
+if __name__ == '__main__':
+    main()
+
+
+# ----------------------------------------------------------------------------

stylegan3-fun/dnnlib/__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.
+
+from .util import EasyDict, make_cache_dir_path

stylegan3-fun/dnnlib/util.py

@@ -0,0 +1,491 @@
+# 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.
+
+"""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
+
+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 format_time_brief(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".format(s // (60 * 60), (s // 60) % 60)
+    else:
+        return "{0}d {1:02}h".format(s // (24 * 60 * 60), (s // (60 * 60)) % 24)
+
+
+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)

stylegan3-fun/docs/configs.md

@@ -0,0 +1,201 @@
+# Training configurations
+
+This document provides guidelines for selecting appropriate training options for various scenarios, as well as an extensive list of recommended configurations.
+
+#### Example
+
+In the remainder of this document, we summarize each configuration as follows:
+
+| <sub>Config</sub><br><br>    | <sub>s/kimg</sub><br><sup>(V100)</sup> | <sub>s/kimg</sub><br><sup>(A100)</sup> | <sub>GPU</sub><br><sup>mem</sup> | <sub>Options</sub><br><br>
+| :--------------------------- | :--------------: | :--------------: | :------------: | :--
+| <sub>StyleGAN3&#8209;T</sub> | <sub>18.47</sub> | <sub>12.29</sub> | <sub>4.3</sub> | <sub>`--cfg=stylegan3-t --gpus=8 --batch=32 --gamma=8.2 --mirror=1`</sub>
+
+This corresponds to the following command line:
+
+```.bash
+# Train StyleGAN3-T for AFHQv2 using 8 GPUs.
+python train.py --outdir=~/training-runs --cfg=stylegan3-t --data=~/datasets/afhqv2-512x512.zip \
+  --gpus=8 --batch=32 --gamma=8.2 --mirror=1
+```
+
+Explanation of the columns:
+- **Config**: StyleGAN3-T (translation equiv.), StyleGAN3-R (translation and rotation equiv.), or StyleGAN2. Reflects the value of `--cfg`.
+- **s/kimg**: Raw training speed, measured separately on Tesla V100 and A100 using our recommended Docker image. The number indicates how many seconds, on average, it takes to process 1000 images from the training set. The number tends to vary slightly over the course of training; typically by no more than &plusmn;20%.
+- **GPU mem**: Maximum GPU memory usage observed during training, reported in gigabytes per GPU. The above example uses 8 GPUs, which means that the total GPU memory usage is around 34.4 GB.
+- **Options**: Command line options for `train.py`, excluding `--outdir` and `--data`.
+
+#### Total training time
+
+In addition the raw s/kimg number, the training time also depends on the `--kimg` and `--metric` options. `--kimg` controls the total number of training iterations and is set to 25000 by default. This is long enough to reach convergence in typical cases, but in practice the results should already look quite reasonable around 5000 kimg. `--metrics` determines which quality metrics are computed periodically during training. The default is `fid50k_full`, which increases the training time slightly; typically by no more than 5%. The automatic computation can be disabled by specifying `--metrics=none`.
+
+In the above example, the total training time on V100 is approximately 18.47 s/kimg * 25000 kimg * 1.05 &thickapprox; 485,000 seconds &thickapprox; 5 days and 14 hours. Disabling metric computation (`--metrics=none`) reduces this to approximately 5 days and 8 hours.
+
+## General guidelines
+
+The most important hyperparameter that needs to be tuned on a per-dataset basis is the R<sub>1</sub> regularization weight, `--gamma`, that must be specified explicitly for `train.py`. As a rule of thumb, the value of `--gamma` scales quadratically with respect to the training set resolution: doubling the resolution (e.g., 256x256 &rarr; 512x512) means that `--gamma` should be multiplied by 4 (e.g., 2 &rarr; 8). The optimal value is usually the same for `--cfg=stylegan3-t` and `--cfg=stylegan3-r`, but considerably lower for `--cfg=stylegan2`.
+
+In practice, we recommend selecting the value of `--gamma` as follows:
+- Find the closest match for your specific case in this document (config, resolution, and GPU count).
+- Try training with the same `--gamma` first.
+- Then, try increasing the value by 2x and 4x, and also decreasing it by 2x and 4x.
+- Pick the value that yields the lowest FID.
+
+The results may also be improved by adjusting `--mirror` and `--aug`, depending on the training data. Specifying `--mirror=1` augments the dataset with random *x*-flips, which effectively doubles the number of images. This is generally beneficial with datasets that are horizontally symmetric (e.g., FFHQ), but it can be harmful if the images contain noticeable asymmetric features (e.g., text or letters). Specifying `--aug=noaug` disables adaptive discriminator augmentation (ADA), which may improve the results slightly if the training set is large enough (at least 100k images when accounting for *x*-flips). With small datasets (less than 30k images), it is generally a good idea to leave the augmentations enabled.
+
+It is possible to speed up the training by decreasing network capacity, i.e., `--cbase=16384`. This typically leads to lower quality results, but the difference is less pronounced with low-resolution datasets (e.g., 256x256).
+
+#### Scaling to different number of GPUs
+
+You can select the number of GPUs by changing the value of `--gpu`; this does not affect the convergence curves or training dynamics in any way. By default, the total batch size (`--batch`) is divided evenly among the GPUs, which means that decreasing the number of GPUs yields higher per-GPU memory usage. To avoid running out of memory, you can decrease the per-GPU batch size by specifying `--batch-gpu`, which performs the same computation in multiple passes using gradient accumulation.
+
+By default, `train.py` exports network snapshots once every 200 kimg, i.e., the product of `--snap=50` and `--tick=4`. When using few GPUs (e.g., 1&ndash;2), this means that it may take a very long time for the first snapshot to appear. We recommend increasing the snapshot frequency in such cases by specifying `--snap=20`, `--snap=10`, or `--snap=5`.
+
+Note that the configurations listed in this document have been specifically tuned for 8 GPUs. The safest way to scale them to different GPU counts is to adjust `--gpu`, `--batch-gpu`, and `--snap` as described above, but it may be possible to reach faster convergence by adjusting some of the other hyperparameters as well. Note, however, that adjusting the total batch size (`--batch`) requires some experimentation; decreasing `--batch` usually necessitates increasing regularization (`--gamma`) and/or decreasing the learning rates (most importantly `--dlr`).
+
+#### Transfer learning
+
+Transfer learning makes it possible to reach very good results very quickly, especially when the training set is small and/or the images resemble the ones produced by a pre-trained model. To enable transfer learning, you can point `--resume` to one of the pre-trained models that we provide for [StyleGAN3](https://ngc.nvidia.com/catalog/models/nvidia:research:stylegan3) and [StyleGAN2](https://ngc.nvidia.com/catalog/models/nvidia:research:stylegan2). For example:
+
+```.bash
+# Fine-tune StyleGAN3-R for MetFaces-U using 1 GPU, starting from the pre-trained FFHQ-U pickle.
+python train.py --outdir=~/training-runs --cfg=stylegan3-r --data=~/datasets/metfacesu-1024x1024.zip \
+    --gpus=8 --batch=32 --gamma=6.6 --mirror=1 --kimg=5000 --snap=5 \
+    --resume=https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-r-ffhqu-1024x1024.pkl
+```
+
+The pre-trained model should be selected to match the specified config, resolution, and architecture-related hyperparameters (e.g., `--cbase`, `--map-depth`, and `--mbstd-group`). You check this by looking at the `fakes_init.png` exported by `train.py` at the beginning; if the configuration is correct, the images should look reasonable.
+
+With transfer learning, the results may be improved slightly by adjusting `--freezed`, in addition to the above guidelines for `--gamma`, `--mirror`, and `--aug`. In our experience, `--freezed=10` and `--freezed=13` tend to work reasonably well.
+
+## Recommended configurations
+
+This section lists recommended settings for StyleGAN3-T and StyleGAN3-R for different resolutions and GPU counts, selected according to the above guidelines. These are intended to provide a good starting point when experimenting with a new dataset. Please note that many of the options (e.g., `--gamma`, `--mirror`, and `--aug`) are still worth adjusting on a case-by-case basis.
+
+#### 128x128 resolution
+
+| <sub>Config</sub><br><br>    | <sub>GPUs</sub><br><br> | <sub>s/kimg</sub><br><sup>(V100)</sup> | <sub>s/kimg</sub><br><sup>(A100)</sup> | <sub>GPU</sub><br><sup>mem</sup> | <sub>Options</sub><br><br>
+| :--------------------------- | :----------: | :--------------: | :--------------: | :------------: | :--
+| <sub>StyleGAN3&#8209;T</sub> | <sub>1</sub> | <sub>73.68</sub> | <sub>27.20</sub> | <sub>7.2</sub> | <sub>`--cfg=stylegan3-t --gpus=1 --batch=32 --gamma=0.5 --batch-gpu=16 --snap=10`</sub>
+| <sub>StyleGAN3&#8209;T</sub> | <sub>2</sub> | <sub>37.30</sub> | <sub>13.74</sub> | <sub>7.1</sub> | <sub>`--cfg=stylegan3-t --gpus=2 --batch=32 --gamma=0.5 --snap=20`</sub>
+| <sub>StyleGAN3&#8209;T</sub> | <sub>4</sub> | <sub>20.66</sub> | <sub>7.52</sub>  | <sub>4.1</sub> | <sub>`--cfg=stylegan3-t --gpus=4 --batch=32 --gamma=0.5`</sub>
+| <sub>StyleGAN3&#8209;T</sub> | <sub>8</sub> | <sub>11.31</sub> | <sub>4.40</sub>  | <sub>2.6</sub> | <sub>`--cfg=stylegan3-t --gpus=8 --batch=32 --gamma=0.5`</sub>
+| <sub>StyleGAN3&#8209;R</sub> | <sub>1</sub> | <sub>58.44</sub> | <sub>34.23</sub> | <sub>8.3</sub> | <sub>`--cfg=stylegan3-r --gpus=1 --batch=32 --gamma=0.5 --batch-gpu=16 --snap=10`</sub>
+| <sub>StyleGAN3&#8209;R</sub> | <sub>2</sub> | <sub>29.92</sub> | <sub>17.29</sub> | <sub>8.2</sub> | <sub>`--cfg=stylegan3-r --gpus=2 --batch=32 --gamma=0.5 --snap=20`</sub>
+| <sub>StyleGAN3&#8209;R</sub> | <sub>4</sub> | <sub>15.49</sub> | <sub>9.53</sub>  | <sub>4.5</sub> | <sub>`--cfg=stylegan3-r --gpus=4 --batch=32 --gamma=0.5`</sub>
+| <sub>StyleGAN3&#8209;R</sub> | <sub>8</sub> | <sub>8.43</sub>  | <sub>5.69</sub>  | <sub>2.7</sub> | <sub>`--cfg=stylegan3-r --gpus=8 --batch=32 --gamma=0.5`</sub>
+
+#### 256x256 resolution
+
+| <sub>Config</sub><br><br>    | <sub>GPUs</sub><br><br> | <sub>s/kimg</sub><br><sup>(V100)</sup> | <sub>s/kimg</sub><br><sup>(A100)</sup> | <sub>GPU</sub><br><sup>mem</sup> | <sub>Options</sub><br><br>
+| :--------------------------- | :----------: | :--------------: | :--------------: | :------------: | :--
+| <sub>StyleGAN3&#8209;T</sub> | <sub>1</sub> | <sub>89.15</sub> | <sub>49.81</sub> | <sub>9.5</sub> | <sub>`--cfg=stylegan3-t --gpus=1 --batch=32 --gamma=2 --batch-gpu=16 --snap=10`</sub>
+| <sub>StyleGAN3&#8209;T</sub> | <sub>2</sub> | <sub>45.45</sub> | <sub>25.05</sub> | <sub>9.3</sub> | <sub>`--cfg=stylegan3-t --gpus=2 --batch=32 --gamma=2 --snap=20`</sub>
+| <sub>StyleGAN3&#8209;T</sub> | <sub>4</sub> | <sub>23.94</sub> | <sub>13.26</sub> | <sub>5.2</sub> | <sub>`--cfg=stylegan3-t --gpus=4 --batch=32 --gamma=2`</sub>
+| <sub>StyleGAN3&#8209;T</sub> | <sub>8</sub> | <sub>13.04</sub> | <sub>7.32</sub>  | <sub>3.1</sub> | <sub>`--cfg=stylegan3-t --gpus=8 --batch=32 --gamma=2`</sub>
+| <sub>StyleGAN3&#8209;R</sub> | <sub>1</sub> | <sub>87.37</sub> | <sub>56.73</sub> | <sub>6.7</sub> | <sub>`--cfg=stylegan3-r --gpus=1 --batch=32 --gamma=2 --batch-gpu=8 --snap=10`</sub>
+| <sub>StyleGAN3&#8209;R</sub> | <sub>2</sub> | <sub>44.12</sub> | <sub>28.60</sub> | <sub>6.7</sub> | <sub>`--cfg=stylegan3-r --gpus=2 --batch=32 --gamma=2 --batch-gpu=8 --snap=20`</sub>
+| <sub>StyleGAN3&#8209;R</sub> | <sub>4</sub> | <sub>22.42</sub> | <sub>14.39</sub> | <sub>6.6</sub> | <sub>`--cfg=stylegan3-r --gpus=4 --batch=32 --gamma=2`</sub>
+| <sub>StyleGAN3&#8209;R</sub> | <sub>8</sub> | <sub>11.88</sub> | <sub>8.03</sub>  | <sub>3.7</sub> | <sub>`--cfg=stylegan3-r --gpus=8 --batch=32 --gamma=2`</sub>
+
+#### 512x512 resolution
+
+| <sub>Config</sub><br><br>    | <sub>GPUs</sub><br><br> | <sub>s/kimg</sub><br><sup>(V100)</sup> | <sub>s/kimg</sub><br><sup>(A100)</sup> | <sub>GPU</sub><br><sup>mem</sup> | <sub>Options</sub><br><br>
+| :--------------------------- | :----------: | :---------------: | :---------------: | :------------: | :--
+| <sub>StyleGAN3&#8209;T</sub> | <sub>1</sub> | <sub>137.33</sub> | <sub>90.25</sub>  | <sub>7.8</sub> | <sub>`--cfg=stylegan3-t --gpus=1 --batch=32 --gamma=8 --batch-gpu=8 --snap=10`</sub>
+| <sub>StyleGAN3&#8209;T</sub> | <sub>2</sub> | <sub>69.65</sub>  | <sub>45.42</sub>  | <sub>7.7</sub> | <sub>`--cfg=stylegan3-t --gpus=2 --batch=32 --gamma=8 --batch-gpu=8 --snap=20`</sub>
+| <sub>StyleGAN3&#8209;T</sub> | <sub>4</sub> | <sub>34.88</sub>  | <sub>22.81</sub>  | <sub>7.6</sub> | <sub>`--cfg=stylegan3-t --gpus=4 --batch=32 --gamma=8`</sub>
+| <sub>StyleGAN3&#8209;T</sub> | <sub>8</sub> | <sub>18.47</sub>  | <sub>12.29</sub>  | <sub>4.3</sub> | <sub>`--cfg=stylegan3-t --gpus=8 --batch=32 --gamma=8`</sub>
+| <sub>StyleGAN3&#8209;R</sub> | <sub>1</sub> | <sub>158.91</sub> | <sub>110.13</sub> | <sub>6.0</sub> | <sub>`--cfg=stylegan3-r --gpus=1 --batch=32 --gamma=8 --batch-gpu=4 --snap=10`</sub>
+| <sub>StyleGAN3&#8209;R</sub> | <sub>2</sub> | <sub>79.96</sub>  | <sub>55.18</sub>  | <sub>6.0</sub> | <sub>`--cfg=stylegan3-r --gpus=2 --batch=32 --gamma=8 --batch-gpu=4 --snap=20`</sub>
+| <sub>StyleGAN3&#8209;R</sub> | <sub>4</sub> | <sub>40.86</sub>  | <sub>27.99</sub>  | <sub>5.9</sub> | <sub>`--cfg=stylegan3-r --gpus=4 --batch=32 --gamma=8 --batch-gpu=4`</sub>
+| <sub>StyleGAN3&#8209;R</sub> | <sub>8</sub> | <sub>20.44</sub>  | <sub>14.04</sub>  | <sub>5.9</sub> | <sub>`--cfg=stylegan3-r --gpus=8 --batch=32 --gamma=8`</sub>
+
+#### 1024x1024 resolution
+
+| <sub>Config</sub><br><br>    | <sub>GPUs</sub><br><br> | <sub>s/kimg</sub><br><sup>(V100)</sup> | <sub>s/kimg</sub><br><sup>(A100)</sup> | <sub>GPU</sub><br><sup>mem</sup> | <sub>Options</sub><br><br>
+| :--------------------------- | :----------: | :---------------: | :---------------: | :-------------: | :--
+| <sub>StyleGAN3&#8209;T</sub> | <sub>1</sub> | <sub>221.85</sub> | <sub>156.91</sub> | <sub>7.0</sub>  | <sub>`--cfg=stylegan3-t --gpus=1 --batch=32 --gamma=32 --batch-gpu=4 --snap=5`</sub>
+| <sub>StyleGAN3&#8209;T</sub> | <sub>2</sub> | <sub>113.44</sub> | <sub>79.16</sub>  | <sub>6.8</sub>  | <sub>`--cfg=stylegan3-t --gpus=2 --batch=32 --gamma=32 --batch-gpu=4 --snap=10`</sub>
+| <sub>StyleGAN3&#8209;T</sub> | <sub>4</sub> | <sub>57.04</sub>  | <sub>39.62</sub>  | <sub>6.7</sub>  | <sub>`--cfg=stylegan3-t --gpus=4 --batch=32 --gamma=32 --batch-gpu=4 --snap=20`</sub>
+| <sub>StyleGAN3&#8209;T</sub> | <sub>8</sub> | <sub>28.71</sub>  | <sub>20.01</sub>  | <sub>6.6</sub>  | <sub>`--cfg=stylegan3-t --gpus=8 --batch=32 --gamma=32`</sub>
+| <sub>StyleGAN3&#8209;R</sub> | <sub>1</sub> | <sub>263.44</sub> | <sub>184.81</sub> | <sub>10.2</sub> | <sub>`--cfg=stylegan3-r --gpus=1 --batch=32 --gamma=32 --batch-gpu=4 --snap=5`</sub>
+| <sub>StyleGAN3&#8209;R</sub> | <sub>2</sub> | <sub>134.22</sub> | <sub>92.58</sub>  | <sub>10.1</sub> | <sub>`--cfg=stylegan3-r --gpus=2 --batch=32 --gamma=32 --batch-gpu=4 --snap=10`</sub>
+| <sub>StyleGAN3&#8209;R</sub> | <sub>4</sub> | <sub>67.33</sub>  | <sub>46.53</sub>  | <sub>10.0</sub> | <sub>`--cfg=stylegan3-r --gpus=4 --batch=32 --gamma=32 --batch-gpu=4 --snap=20`</sub>
+| <sub>StyleGAN3&#8209;R</sub> | <sub>8</sub> | <sub>34.12</sub>  | <sub>23.42</sub>  | <sub>9.9</sub>  | <sub>`--cfg=stylegan3-r --gpus=8 --batch=32 --gamma=32`</sub>
+
+## Configurations used in StyleGAN3 paper
+
+This section lists the exact settings that we used in the "Alias-Free Generative Adversarial Networks" paper.
+
+#### FFHQ-U and FFHQ at 1024x1024 resolution
+
+| <sub>Config</sub><br><br>    | <sub>s/kimg</sub><br><sup>(V100)</sup> | <sub>s/kimg</sub><br><sup>(A100)</sup> | <sub>GPU</sub><br><sup>mem</sup> | <sub>Options</sub><br><br>
+| :--------------------------- | :--------------: | :--------------: | :------------: | :--
+| <sub>StyleGAN2</sub>         | <sub>17.55</sub> | <sub>14.57</sub> | <sub>6.2</sub> | <sub>`--cfg=stylegan2 --gpus=8 --batch=32 --gamma=10 --mirror=1 --aug=noaug`</sub>
+| <sub>StyleGAN3&#8209;T</sub> | <sub>28.71</sub> | <sub>20.01</sub> | <sub>6.6</sub> | <sub>`--cfg=stylegan3-t --gpus=8 --batch=32 --gamma=32.8 --mirror=1 --aug=noaug`</sub>
+| <sub>StyleGAN3&#8209;R</sub> | <sub>34.12</sub> | <sub>23.42</sub> | <sub>9.9</sub> | <sub>`--cfg=stylegan3-r --gpus=8 --batch=32 --gamma=32.8 --mirror=1 --aug=noaug`</sub>
+
+#### MetFaces-U at 1024x1024 resolution
+
+| <sub>Config</sub><br><br>    | <sub>s/kimg</sub><br><sup>(V100)</sup> | <sub>s/kimg</sub><br><sup>(A100)</sup> | <sub>GPU</sub><br><sup>mem</sup> | <sub>Options</sub><br><br>
+| :--------------------------- | :--------------: | :--------------: | :-------------: | :--
+| <sub>StyleGAN2</sub>         | <sub>18.74</sub> | <sub>11.80</sub> | <sub>7.4</sub>  | <sub>`--cfg=stylegan2 --gpus=8 --batch=32 --gamma=10 --mirror=1 --kimg=5000 --snap=10 --resume=https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan2/versions/1/files/stylegan2-ffhqu-1024x1024.pkl`</sub>
+| <sub>StyleGAN3&#8209;T</sub> | <sub>29.84</sub> | <sub>21.06</sub> | <sub>7.7</sub>  | <sub>`--cfg=stylegan3-t --gpus=8 --batch=32 --gamma=16.4 --mirror=1 --kimg=5000 --snap=10 --resume=https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-t-ffhqu-1024x1024.pkl`</sub>
+| <sub>StyleGAN3&#8209;R</sub> | <sub>35.10</sub> | <sub>24.32</sub> | <sub>10.9</sub> | <sub>`--cfg=stylegan3-r --gpus=8 --batch=32 --gamma=6.6 --mirror=1 --kimg=5000 --snap=10 --resume=https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-r-ffhqu-1024x1024.pkl`</sub>
+
+#### MetFaces at 1024x1024 resolution
+
+| <sub>Config</sub><br><br>    | <sub>s/kimg</sub><br><sup>(V100)</sup> | <sub>s/kimg</sub><br><sup>(A100)</sup> | <sub>GPU</sub><br><sup>mem</sup> | <sub>Options</sub><br><br>
+| :--------------------------- | :--------------: | :--------------: | :-------------: | :--
+| <sub>StyleGAN2</sub>         | <sub>18.74</sub> | <sub>11.80</sub> | <sub>7.4</sub>  | <sub>`--cfg=stylegan2 --gpus=8 --batch=32 --gamma=5 --mirror=1 --kimg=5000 --snap=10 --resume=https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan2/versions/1/files/stylegan2-ffhq-1024x1024.pkl`</sub>
+| <sub>StyleGAN3&#8209;T</sub> | <sub>29.84</sub> | <sub>21.06</sub> | <sub>7.7</sub>  | <sub>`--cfg=stylegan3-t --gpus=8 --batch=32 --gamma=6.6 --mirror=1 --kimg=5000 --snap=10 --resume=https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-t-ffhq-1024x1024.pkl`</sub>
+| <sub>StyleGAN3&#8209;R</sub> | <sub>35.10</sub> | <sub>24.32</sub> | <sub>10.9</sub> | <sub>`--cfg=stylegan3-r --gpus=8 --batch=32 --gamma=3.3 --mirror=1 --kimg=5000 --snap=10 --resume=https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-r-ffhq-1024x1024.pkl`</sub>
+
+#### AFHQv2 at 512x512 resolution
+
+| <sub>Config</sub><br><br>    | <sub>s/kimg</sub><br><sup>(V100)</sup> | <sub>s/kimg</sub><br><sup>(A100)</sup> | <sub>GPU</sub><br><sup>mem</sup> | <sub>Options</sub><br><br>
+| :--------------------------- | :--------------: | :--------------: | :------------: | :--
+| <sub>StyleGAN2</sub>         | <sub>10.90</sub> | <sub>6.60</sub>  | <sub>3.9</sub> | <sub>`--cfg=stylegan2 --gpus=8 --batch=32 --gamma=5 --mirror=1`</sub>
+| <sub>StyleGAN3&#8209;T</sub> | <sub>18.47</sub> | <sub>12.29</sub> | <sub>4.3</sub> | <sub>`--cfg=stylegan3-t --gpus=8 --batch=32 --gamma=8.2 --mirror=1`</sub>
+| <sub>StyleGAN3&#8209;R</sub> | <sub>20.44</sub> | <sub>14.04</sub> | <sub>5.9</sub> | <sub>`--cfg=stylegan3-r --gpus=8 --batch=32 --gamma=16.4 --mirror=1`</sub>
+
+#### FFHQ-U ablations at 256x256 resolution
+
+| <sub>Config</sub><br><br>    | <sub>s/kimg</sub><br><sup>(V100)</sup> | <sub>s/kimg</sub><br><sup>(A100)</sup> | <sub>GPU</sub><br><sup>mem</sup> | <sub>Options</sub><br><br>
+| :--------------------------- | :-------------: | :-------------: | :------------: | :--
+| <sub>StyleGAN2</sub>         | <sub>3.61</sub> | <sub>2.19</sub> | <sub>2.7</sub> | <sub>`--cfg=stylegan2 --gpus=8 --batch=64 --gamma=1 --mirror=1 --aug=noaug --cbase=16384 --glr=0.0025 --dlr=0.0025 --mbstd-group=8`</sub>
+| <sub>StyleGAN3&#8209;T</sub> | <sub>7.40</sub> | <sub>3.74</sub> | <sub>3.5</sub> | <sub>`--cfg=stylegan3-t --gpus=8 --batch=64 --gamma=1 --mirror=1 --aug=noaug --cbase=16384 --dlr=0.0025`</sub>
+| <sub>StyleGAN3&#8209;R</sub> | <sub>6.71</sub> | <sub>4.81</sub> | <sub>4.2</sub> | <sub>`--cfg=stylegan3-r --gpus=8 --batch=64 --gamma=1 --mirror=1 --aug=noaug --cbase=16384 --dlr=0.0025`</sub>
+
+## Old StyleGAN2-ADA configurations
+
+This section lists command lines that can be used to match the configurations provided by our previous [StyleGAN2-ADA](https://github.com/NVlabs/stylegan2-ada-pytorch) codebase. The first table corresponds to `--cfg=auto` (default) for different resolutions and GPU counts, while the second table lists the remaining alternatives.
+
+#### Default configuration
+
+| <sub>Res.</sub><br><br> | <sub>GPUs</sub><br><br> | <sub>s/kimg</sub><br><sup>(V100)</sup> | <sub>s/kimg</sub><br><sup>(A100)</sup> | <sub>GPU</sub><br><sup>mem</sup> | <sub>Options</sub><br><br>
+| :---------------------- | :----------: | :---------------: | :--------------: | :------------: | :--
+| <sub>128&sup2;</sub>    | <sub>1</sub> | <sub>12.51</sub>  | <sub>6.79</sub>  | <sub>6.2</sub> | <sub>`--cfg=stylegan2 --gpus=1 --batch=32 --gamma=0.1024 --map-depth=2 --glr=0.0025 --dlr=0.0025 --cbase=16384`</sub>
+| <sub>128&sup2;</sub>    | <sub>2</sub> | <sub>6.43</sub>   | <sub>3.45</sub>  | <sub>6.2</sub> | <sub>`--cfg=stylegan2 --gpus=2 --batch=64 --gamma=0.0512 --map-depth=2 --glr=0.0025 --dlr=0.0025 --cbase=16384`</sub>
+| <sub>128&sup2;</sub>    | <sub>4</sub> | <sub>3.82</sub>   | <sub>2.23</sub>  | <sub>3.5</sub> | <sub>`--cfg=stylegan2 --gpus=4 --batch=64 --gamma=0.0512 --map-depth=2 --glr=0.0025 --dlr=0.0025 --cbase=16384`</sub>
+| <sub>256&sup2;</sub>    | <sub>1</sub> | <sub>20.84</sub>  | <sub>12.53</sub> | <sub>4.5</sub> | <sub>`--cfg=stylegan2 --gpus=1 --batch=16 --gamma=0.8192 --map-depth=2 --glr=0.0025 --dlr=0.0025 --cbase=16384`</sub>
+| <sub>256&sup2;</sub>    | <sub>2</sub> | <sub>10.93</sub>  | <sub>6.36</sub>  | <sub>4.5</sub> | <sub>`--cfg=stylegan2 --gpus=2 --batch=32 --gamma=0.4096 --map-depth=2 --glr=0.0025 --dlr=0.0025 --cbase=16384`</sub>
+| <sub>256&sup2;</sub>    | <sub>4</sub> | <sub>5.39</sub>   | <sub>3.20</sub>  | <sub>4.5</sub> | <sub>`--cfg=stylegan2 --gpus=4 --batch=64 --gamma=0.2048 --map-depth=2 --glr=0.0025 --dlr=0.0025 --cbase=16384`</sub>
+| <sub>256&sup2;</sub>    | <sub>8</sub> | <sub>3.89</sub>   | <sub>2.38</sub>  | <sub>2.6</sub> | <sub>`--cfg=stylegan2 --gpus=8 --batch=64 --gamma=0.2048 --map-depth=2 --glr=0.0025 --dlr=0.0025 --cbase=16384`</sub>
+| <sub>512&sup2;</sub>    | <sub>1</sub> | <sub>71.59</sub>  | <sub>41.06</sub> | <sub>6.8</sub> | <sub>`--cfg=stylegan2 --gpus=1 --batch=8 --gamma=6.5536 --map-depth=2 --glr=0.0025 --dlr=0.0025`</sub>
+| <sub>512&sup2;</sub>    | <sub>2</sub> | <sub>36.79</sub>  | <sub>20.83</sub> | <sub>6.8</sub> | <sub>`--cfg=stylegan2 --gpus=2 --batch=16 --gamma=3.2768 --map-depth=2 --glr=0.0025 --dlr=0.0025`</sub>
+| <sub>512&sup2;</sub>    | <sub>4</sub> | <sub>18.12</sub>  | <sub>10.45</sub> | <sub>6.7</sub> | <sub>`--cfg=stylegan2 --gpus=4 --batch=32 --gamma=1.6384 --map-depth=2 --glr=0.0025 --dlr=0.0025`</sub>
+| <sub>512&sup2;</sub>    | <sub>8</sub> | <sub>9.09</sub>   | <sub>5.24</sub>  | <sub>6.8</sub> | <sub>`--cfg=stylegan2 --gpus=8 --batch=64 --gamma=0.8192 --map-depth=2 --glr=0.0025 --dlr=0.0025`</sub>
+| <sub>1024&sup2;</sub>   | <sub>1</sub> | <sub>141.83</sub> | <sub>90.39</sub> | <sub>7.2</sub> | <sub>`--cfg=stylegan2 --gpus=1 --batch=4 --gamma=52.4288 --map-depth=2`</sub>
+| <sub>1024&sup2;</sub>   | <sub>2</sub> | <sub>73.13</sub>  | <sub>46.04</sub> | <sub>7.2</sub> | <sub>`--cfg=stylegan2 --gpus=2 --batch=8 --gamma=26.2144 --map-depth=2`</sub>
+| <sub>1024&sup2;</sub>   | <sub>4</sub> | <sub>36.95</sub>  | <sub>23.15</sub> | <sub>7.0</sub> | <sub>`--cfg=stylegan2 --gpus=4 --batch=16 --gamma=13.1072 --map-depth=2`</sub>
+| <sub>1024&sup2;</sub>   | <sub>8</sub> | <sub>18.47</sub>  | <sub>11.66</sub> | <sub>7.3</sub> | <sub>`--cfg=stylegan2 --gpus=8 --batch=32 --gamma=6.5536 --map-depth=2`</sub>
+
+#### Repro configurations
+
+| <sub>Name</sub><br><br> | <sub>s/kimg</sub><br><sup>(V100)</sup> | <sub>s/kimg</sub><br><sup>(A100)</sup> | <sub>GPU</sub><br><sup>mem</sup> | <sub>Options</sub><br><br>
+| :---------------------- | :--------------: | :--------------: | :------------: | :--
+| <sub>`stylegan2`</sub>  | <sub>17.55</sub> | <sub>14.57</sub> | <sub>6.2</sub> | <sub>`--cfg=stylegan2 --gpus=8 --batch=32 --gamma=10`</sub>
+| <sub>`paper256`</sub>   | <sub>4.01</sub>  | <sub>2.47</sub>  | <sub>2.7</sub> | <sub>`--cfg=stylegan2 --gpus=8 --batch=64 --gamma=1 --cbase=16384 --glr=0.0025 --dlr=0.0025 --mbstd-group=8`</sub>
+| <sub>`paper512`</sub>   | <sub>9.11</sub>  | <sub>5.28</sub>  | <sub>6.7</sub> | <sub>`--cfg=stylegan2 --gpus=8 --batch=64 --gamma=0.5 --glr=0.0025 --dlr=0.0025 --mbstd-group=8`</sub>
+| <sub>`paper1024`</sub>  | <sub>18.56</sub> | <sub>11.75</sub> | <sub>6.9</sub> | <sub>`--cfg=stylegan2 --gpus=8 --batch=32 --gamma=2`</sub>

stylegan3-fun/docs/dataset-tool-help.txt

@@ -0,0 +1,70 @@
+Usage: dataset_tool.py [OPTIONS]
+
+  Convert an image dataset into a dataset archive usable with StyleGAN2 ADA
+  PyTorch.
+
+  The input dataset format is guessed from the --source argument:
+
+  --source *_lmdb/                    Load LSUN dataset
+  --source cifar-10-python.tar.gz     Load CIFAR-10 dataset
+  --source train-images-idx3-ubyte.gz Load MNIST dataset
+  --source path/                      Recursively load all images from path/
+  --source dataset.zip                Recursively load all images from dataset.zip
+
+  Specifying the output format and path:
+
+  --dest /path/to/dir                 Save output files under /path/to/dir
+  --dest /path/to/dataset.zip         Save output files into /path/to/dataset.zip
+
+  The output dataset format can be either an image folder or an uncompressed
+  zip archive. Zip archives makes it easier to move datasets around file
+  servers and clusters, and may offer better training performance on network
+  file systems.
+
+  Images within the dataset archive will be stored as uncompressed PNG.
+  Uncompresed PNGs can be efficiently decoded in the training loop.
+
+  Class labels are stored in a file called 'dataset.json' that is stored at
+  the dataset root folder.  This file has the following structure:
+
+  {
+      "labels": [
+          ["00000/img00000000.png",6],
+          ["00000/img00000001.png",9],
+          ... repeated for every image in the datase
+          ["00049/img00049999.png",1]
+      ]
+  }
+
+  If the 'dataset.json' file cannot be found, the dataset is interpreted as
+  not containing class labels.
+
+  Image scale/crop and resolution requirements:
+
+  Output images must be square-shaped and they must all have the same power-
+  of-two dimensions.
+
+  To scale arbitrary input image size to a specific width and height, use
+  the --resolution option.  Output resolution will be either the original
+  input resolution (if resolution was not specified) or the one specified
+  with --resolution option.
+
+  Use the --transform=center-crop or --transform=center-crop-wide options to
+  apply a center crop transform on the input image.  These options should be
+  used with the --resolution option.  For example:
+
+  python dataset_tool.py --source LSUN/raw/cat_lmdb --dest /tmp/lsun_cat \
+      --transform=center-crop-wide --resolution=512x384
+
+Options:
+  --source PATH                   Directory or archive name for input dataset
+                                  [required]
+
+  --dest PATH                     Output directory or archive name for output
+                                  dataset  [required]
+
+  --max-images INTEGER            Output only up to `max-images` images
+  --transform [center-crop|center-crop-wide]
+                                  Input crop/resize mode
+  --resolution WxH                Output resolution (e.g., '512x512')
+  --help                          Show this message and exit.

stylegan3-fun/docs/train-help.txt

@@ -0,0 +1,53 @@
+Usage: train.py [OPTIONS]
+
+  Train a GAN using the techniques described in the paper "Alias-Free
+  Generative Adversarial Networks".
+
+  Examples:
+
+  # Train StyleGAN3-T for AFHQv2 using 8 GPUs.
+  python train.py --outdir=~/training-runs --cfg=stylegan3-t --data=~/datasets/afhqv2-512x512.zip \
+      --gpus=8 --batch=32 --gamma=8.2 --mirror=1
+
+  # Fine-tune StyleGAN3-R for MetFaces-U using 1 GPU, starting from the pre-trained FFHQ-U pickle.
+  python train.py --outdir=~/training-runs --cfg=stylegan3-r --data=~/datasets/metfacesu-1024x1024.zip \
+      --gpus=8 --batch=32 --gamma=6.6 --mirror=1 --kimg=5000 --snap=5 \
+      --resume=https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-r-ffhqu-1024x1024.pkl
+
+  # Train StyleGAN2 for FFHQ at 1024x1024 resolution using 8 GPUs.
+  python train.py --outdir=~/training-runs --cfg=stylegan2 --data=~/datasets/ffhq-1024x1024.zip \
+      --gpus=8 --batch=32 --gamma=10 --mirror=1 --aug=noaug
+
+Options:
+  --outdir DIR                    Where to save the results  [required]
+  --cfg [stylegan3-t|stylegan3-r|stylegan2]
+                                  Base configuration  [required]
+  --data [ZIP|DIR]                Training data  [required]
+  --gpus INT                      Number of GPUs to use  [required]
+  --batch INT                     Total batch size  [required]
+  --gamma FLOAT                   R1 regularization weight  [required]
+  --cond BOOL                     Train conditional model  [default: False]
+  --mirror BOOL                   Enable dataset x-flips  [default: False]
+  --aug [noaug|ada|fixed]         Augmentation mode  [default: ada]
+  --resume [PATH|URL]             Resume from given network pickle
+  --freezed INT                   Freeze first layers of D  [default: 0]
+  --p FLOAT                       Probability for --aug=fixed  [default: 0.2]
+  --target FLOAT                  Target value for --aug=ada  [default: 0.6]
+  --batch-gpu INT                 Limit batch size per GPU
+  --cbase INT                     Capacity multiplier  [default: 32768]
+  --cmax INT                      Max. feature maps  [default: 512]
+  --glr FLOAT                     G learning rate  [default: varies]
+  --dlr FLOAT                     D learning rate  [default: 0.002]
+  --map-depth INT                 Mapping network depth  [default: varies]
+  --mbstd-group INT               Minibatch std group size  [default: 4]
+  --desc STR                      String to include in result dir name
+  --metrics [NAME|A,B,C|none]     Quality metrics  [default: fid50k_full]
+  --kimg KIMG                     Total training duration  [default: 25000]
+  --tick KIMG                     How often to print progress  [default: 4]
+  --snap TICKS                    How often to save snapshots  [default: 50]
+  --seed INT                      Random seed  [default: 0]
+  --fp32 BOOL                     Disable mixed-precision  [default: False]
+  --nobench BOOL                  Disable cuDNN benchmarking  [default: False]
+  --workers INT                   DataLoader worker processes  [default: 3]
+  -n, --dry-run                   Print training options and exit
+  --help                          Show this message and exit.

stylegan3-fun/docs/troubleshooting.md

@@ -0,0 +1,31 @@
+# Troubleshooting
+
+Our PyTorch code uses custom [CUDA extensions](https://pytorch.org/tutorials/advanced/cpp_extension.html) to speed up some of the network layers.  Getting these to run can sometimes be a hassle.
+
+This page aims to give guidance on how to diagnose and fix run-time problems related to these extensions.
+
+## Before you start
+
+1. Try Docker first!  Ensure you can successfully run our models using the recommended Docker image.  Follow the instructions in [README.md](/README.md) to get it running.
+2. Can't use Docker?  Read on..
+
+## Installing dependencies
+
+Make sure you've installed everything listed on the requirements section in the [README.md](/README.md).  The key components w.r.t. custom extensions are:
+
+- **[CUDA toolkit 11.1](https://developer.nvidia.com/cuda-toolkit)** or later (this is not the same as `cudatoolkit` from Conda).
+  - PyTorch invokes `nvcc` to compile our CUDA kernels.
+- **ninja**
+  - PyTorch uses [Ninja](https://ninja-build.org/) as its build system.
+- **GCC** (Linux) or **Visual Studio** (Windows)
+  - GCC 7.x or later is required.  Earlier versions such as GCC 6.3 [are known not to work](https://github.com/NVlabs/stylegan3/issues/2).
+
+#### Why is CUDA toolkit installation necessary?
+
+The PyTorch package contains the required CUDA toolkit libraries needed to run PyTorch, so why is a separate CUDA toolkit installation required?  Our models use custom CUDA kernels to implement operations such as efficient resampling of 2D images.  PyTorch code invokes the CUDA compiler at run-time to compile these kernels on first-use.  The tools and libraries required for this compilation are not bundled in PyTorch and thus a host CUDA toolkit installation is required.
+
+## Things to try
+
+- Completely remove: `$HOME/.cache/torch_extensions` (Linux) or `C:\Users\<username>\AppData\Local\torch_extensions\torch_extensions\Cache` (Windows) and re-run StyleGAN3 python code.
+- Run ninja in `$HOME/.cache/torch_extensions` to see that it builds.
+- Inspect the `build.ninja` in the build directories under `$HOME/.cache/torch_extensions` and check CUDA tools and versions are consistent with what you intended to use.

stylegan3-fun/environment.yml

@@ -0,0 +1,35 @@
+name: stylegan3
+channels:
+  - pytorch
+  - nvidia
+  - conda-forge  # PR #80 by @SetZero / #143 by @coldwaterq
+dependencies:
+  - python >= 3.8
+  - pip
+  - numpy>=1.20
+  - click>=8.0
+  - pillow=8.3.1
+  - scipy=1.7.1
+  - pytorch::pytorch=1.11  # We can now use Pytorch 1.11!
+  - nvidia::cudatoolkit=11.3  # Necessary to remove previous PR by @edstoica
+  - requests=2.26.0
+  - tqdm=4.62.2
+  - ninja=1.10.2
+  - matplotlib=3.4.2
+  - imageio=2.9.0
+  - pip:
+    - imgui==1.3.0
+    - glfw==2.2.0
+    - pyopengl==3.1.5
+    - imageio-ffmpeg==0.4.3
+    - pyspng
+    - pyperlin  # for Discriminator Synthesis
+    - psutil  # PR #125 by @fastflair / #111 by @siddharthksah
+    - tensorboard  # PR #125 by @fastflair
+    - torchvision==0.12.0  # For "Discriminator Synthesis" / discriminator_synthesis.py
+    - pyperlin  # For "Discriminator Synthesis" / discriminator_synthesis.py
+    - scikit-learn  # For "Self-Distilled StyleGAN" / multimodal_truncation.py
+    - moviepy==1.0.3
+    - ffmpeg-python==0.2.0
+    - scikit-video==1.1.11
+    - setuptools==59.5.0   # PR #8 by @ZibbeZabbe

stylegan3-fun/gen_images.py

@@ -0,0 +1,145 @@
+# 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.
+
+"""Generate images using pretrained network pickle."""
+
+import os
+import re
+from typing import List, Optional, Tuple, Union
+
+import click
+import dnnlib
+import numpy as np
+import PIL.Image
+import torch
+
+import legacy
+
+#----------------------------------------------------------------------------
+
+def parse_range(s: Union[str, List]) -> List[int]:
+    '''Parse a comma separated list of numbers or ranges and return a list of ints.
+
+    Example: '1,2,5-10' returns [1, 2, 5, 6, 7]
+    '''
+    if isinstance(s, list): return s
+    ranges = []
+    range_re = re.compile(r'^(\d+)-(\d+)$')
+    for p in s.split(','):
+        m = range_re.match(p)
+        if m:
+            ranges.extend(range(int(m.group(1)), int(m.group(2))+1))
+        else:
+            ranges.append(int(p))
+    return ranges
+
+#----------------------------------------------------------------------------
+
+def parse_vec2(s: Union[str, Tuple[float, float]]) -> Tuple[float, float]:
+    '''Parse a floating point 2-vector of syntax 'a,b'.
+
+    Example:
+        '0,1' returns (0,1)
+    '''
+    if isinstance(s, tuple): return s
+    parts = s.split(',')
+    if len(parts) == 2:
+        return (float(parts[0]), float(parts[1]))
+    raise ValueError(f'cannot parse 2-vector {s}')
+
+#----------------------------------------------------------------------------
+
+def make_transform(translate: Tuple[float,float], angle: float):
+    m = np.eye(3)
+    s = np.sin(angle/360.0*np.pi*2)
+    c = np.cos(angle/360.0*np.pi*2)
+    m[0][0] = c
+    m[0][1] = s
+    m[0][2] = translate[0]
+    m[1][0] = -s
+    m[1][1] = c
+    m[1][2] = translate[1]
+    return m
+
+#----------------------------------------------------------------------------
+
+@click.command()
+@click.option('--network', 'network_pkl', help='Network pickle filename', required=True)
+@click.option('--seeds', type=parse_range, help='List of random seeds (e.g., \'0,1,4-6\')', required=True)
+@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('--translate', help='Translate XY-coordinate (e.g. \'0.3,1\')', type=parse_vec2, default='0,0', show_default=True, metavar='VEC2')
+@click.option('--rotate', help='Rotation angle in degrees', type=float, default=0, show_default=True, metavar='ANGLE')
+@click.option('--outdir', help='Where to save the output images', type=str, required=True, metavar='DIR')
+def generate_images(
+    network_pkl: str,
+    seeds: List[int],
+    truncation_psi: float,
+    noise_mode: str,
+    outdir: str,
+    translate: Tuple[float,float],
+    rotate: float,
+    class_idx: Optional[int]
+):
+    """Generate images using pretrained network pickle.
+
+    Examples:
+
+    \b
+    # Generate an image using pre-trained AFHQv2 model ("Ours" in Figure 1, left).
+    python gen_images.py --outdir=out --trunc=1 --seeds=2 \\
+        --network=https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-r-afhqv2-512x512.pkl
+
+    \b
+    # Generate uncurated images with truncation using the MetFaces-U dataset
+    python gen_images.py --outdir=out --trunc=0.7 --seeds=600-605 \\
+        --network=https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-t-metfacesu-1024x1024.pkl
+    """
+
+    print('Loading networks from "%s"...' % network_pkl)
+    device = torch.device('cuda')
+    with dnnlib.util.open_url(network_pkl) as f:
+        G = legacy.load_network_pkl(f)['G_ema'].to(device) # type: ignore
+
+    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:
+            raise click.ClickException('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')
+
+    # Generate images.
+    for seed_idx, seed in enumerate(seeds):
+        print('Generating image for seed %d (%d/%d) ...' % (seed, seed_idx, len(seeds)))
+        z = torch.from_numpy(np.random.RandomState(seed).randn(1, G.z_dim)).to(device)
+
+        # Construct an inverse rotation/translation matrix and pass to the generator.  The
+        # generator expects this matrix as an inverse to avoid potentially failing numerical
+        # operations in the network.
+        if hasattr(G.synthesis, 'input'):
+            m = make_transform(translate, rotate)
+            m = np.linalg.inv(m)
+            G.synthesis.input.transform.copy_(torch.from_numpy(m))
+
+        img = G(z, label, truncation_psi=truncation_psi, noise_mode=noise_mode)
+        img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)
+        PIL.Image.fromarray(img[0].cpu().numpy(), 'RGB').save(f'{outdir}/seed{seed:04d}.png')
+
+
+#----------------------------------------------------------------------------
+
+if __name__ == "__main__":
+    generate_images() # pylint: disable=no-value-for-parameter
+
+#----------------------------------------------------------------------------

stylegan3-fun/gen_video.py

@@ -0,0 +1,281 @@
+# 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.
+
+"""Generate lerp videos using pretrained network pickle."""
+
+import copy
+import os
+import re
+from typing import List, Optional, Tuple, Union
+
+import click
+import dnnlib
+import imageio
+import numpy as np
+import scipy.interpolate
+import torch
+from tqdm import tqdm
+
+import legacy
+from torch_utils import gen_utils
+
+
+# ----------------------------------------------------------------------------
+
+
+def layout_grid(img, grid_w=None, grid_h=1, float_to_uint8=True, chw_to_hwc=True, to_numpy=True):
+    batch_size, channels, img_h, img_w = img.shape
+    if grid_w is None:
+        grid_w = batch_size // grid_h
+    assert batch_size == grid_w * grid_h
+    if float_to_uint8:
+        img = (img * 127.5 + 128).clamp(0, 255).to(torch.uint8)
+    img = img.reshape(grid_h, grid_w, channels, img_h, img_w)
+    img = img.permute(2, 0, 3, 1, 4)
+    img = img.reshape(channels, grid_h * img_h, grid_w * img_w)
+    if chw_to_hwc:
+        img = img.permute(1, 2, 0)
+    if to_numpy:
+        img = img.cpu().numpy()
+    return img
+
+
+# ----------------------------------------------------------------------------
+
+
+def parse_vec2(s: Union[str, Tuple[float, float]]) -> Tuple[float, float]:
+    """Parse a floating point 2-vector of syntax 'a,b'.
+
+    Example:
+        '0,1' returns (0,1)
+    """
+    if isinstance(s, tuple): return s
+    parts = s.split(',')
+    if len(parts) == 2:
+        return (float(parts[0]), float(parts[1]))
+    raise ValueError(f'cannot parse 2-vector {s}')
+
+
+# ----------------------------------------------------------------------------
+
+
+def make_transform(translate: Tuple[float,float], angle: float):
+    m = np.eye(3)
+    s = np.sin(angle/360.0*np.pi*2)
+    c = np.cos(angle/360.0*np.pi*2)
+    m[0][0] = c
+    m[0][1] = s
+    m[0][2] = translate[0]
+    m[1][0] = -s
+    m[1][1] = c
+    m[1][2] = translate[1]
+    return m
+
+
+# ----------------------------------------------------------------------------
+
+
+def gen_interp_video(G,
+                     mp4: str,
+                     seeds: List[int],
+                     shuffle_seed: int = None,
+                     w_frames: int = 60*4,
+                     kind: str = 'cubic',
+                     grid_dims: Tuple[int] = (1,1),
+                     num_keyframes: int = None,
+                     wraps: int = 2,
+                     psi: float = 1.0,
+                     device: torch.device = torch.device('cuda'),
+                     stabilize_video: bool = True,
+                     **video_kwargs):
+    grid_w = grid_dims[0]
+    grid_h = grid_dims[1]
+
+    if stabilize_video:
+        # Thanks to @RiversHaveWings and @nshepperd1
+        if hasattr(G.synthesis, 'input'):
+            shift = G.synthesis.input.affine(G.mapping.w_avg.unsqueeze(0))
+            G.synthesis.input.affine.bias.data.add_(shift.squeeze(0))
+            G.synthesis.input.affine.weight.data.zero_()
+
+    # Get the Generator's transform
+    m = G.synthesis.input.transform if hasattr(G.synthesis, 'input') else None
+
+    if num_keyframes is None:
+        if len(seeds) % (grid_w*grid_h) != 0:
+            raise ValueError('Number of input seeds must be divisible by grid W*H')
+        num_keyframes = len(seeds) // (grid_w*grid_h)
+
+    all_seeds = np.zeros(num_keyframes*grid_h*grid_w, dtype=np.int64)
+    for idx in range(num_keyframes*grid_h*grid_w):
+        all_seeds[idx] = seeds[idx % len(seeds)]
+
+    if shuffle_seed is not None:
+        rng = np.random.RandomState(seed=shuffle_seed)
+        rng.shuffle(all_seeds)
+
+    zs = torch.from_numpy(np.stack([np.random.RandomState(seed).randn(G.z_dim) for seed in all_seeds])).to(device)
+    ws = G.mapping(z=zs, c=None, truncation_psi=psi)
+    _ = G.synthesis(ws[:1]) # warm up
+    ws = ws.reshape(grid_h, grid_w, num_keyframes, *ws.shape[1:])
+
+    # Interpolation.
+    grid = []
+    for yi in range(grid_h):
+        row = []
+        for xi in range(grid_w):
+            x = np.arange(-num_keyframes * wraps, num_keyframes * (wraps + 1))
+            y = np.tile(ws[yi][xi].cpu().numpy(), [wraps * 2 + 1, 1, 1])
+            interp = scipy.interpolate.interp1d(x, y, kind=kind, axis=0)
+            row.append(interp)
+        grid.append(row)
+
+    # Render video.
+    video_out = imageio.get_writer(mp4, mode='I', fps=60, codec='libx264', **video_kwargs)
+    for frame_idx in tqdm(range(num_keyframes * w_frames)):
+        imgs = []
+        # Construct an inverse affine matrix and pass to the generator. The generator expects 
+        # this matrix as an inverse to avoid potentially failing numerical operations in the network.
+        if hasattr(G.synthesis, 'input'):
+            # Set default values for each affine transformation
+            total_rotation = 0.0  # If >= 0.0, will rotate the pixels counter-clockwise w.r.t. the center; in radians
+            total_translation_x = 0.0  # If >= 0.0, will translate all pixels to the right; if <= 0.0, to the left
+            total_translation_y = 0.0  # If >= 0.0, will translate all pixels upwards; if <= 0.0, downwards
+            total_scale_x = 1.0  # If <= 1.0, will zoom in; else, will zoom out (x-axis)
+            total_scale_y = 1.0  # If <= 1.0, will zoom in; else, will zoom out (y-axis)
+            total_shear_x = 0.0  # If >= 0.0, will shear pixels to the right, keeping y fixed; if <= 0.0, to the left
+            total_shear_y = 0.0  # If >= 0.0, will shear pixels upwards, keeping x fixed; if <= 0.0, downwards
+            mirror_x = False  # Mirror along the x-axis; if True, will flip the image horizontally (can't be a function of frame_idx)
+            mirror_y = False  # Mirror along the y-axis; if True, will flip the image vertically (can't be a function of frame_idx)
+
+            # Go nuts with these. They can be constants as above to fix centering/rotation in your video, 
+            # or you can make them functions of frame_idx to animate them, such as (uncomment as many as you want to try):
+            # total_scale_x = 1 + np.sin(np.pi*frame_idx/(num_keyframes * w_frames))/2  # will oscillate between 0.5 and 1.5
+            # total_rotation = 4*np.pi*frame_idx/(num_keyframes * w_frames)  # 4 will dictate the number of rotations, so 1 full rotation
+            # total_shear_y = 2*np.sin(2*np.pi*frame_idx/(num_keyframes * w_frames))  # will oscillate between -2 and 2
+
+            # We then use these values to construct the affine matrix
+            m = gen_utils.make_affine_transform(m, angle=total_rotation, translate_x=total_translation_x,
+                                                translate_y=total_translation_y, scale_x=total_scale_x,
+                                                scale_y=total_scale_y, shear_x=total_shear_x, shear_y=total_shear_y,
+                                                mirror_x=mirror_x, mirror_y=mirror_y)
+            m = np.linalg.inv(m)
+            # Finally, we pass the matrix to the generator
+            G.synthesis.input.transform.copy_(torch.from_numpy(m))
+
+        # The rest stays the same, for all you gen_video.py lovers out there
+        for yi in range(grid_h):
+            for xi in range(grid_w):
+                interp = grid[yi][xi]
+                w = torch.from_numpy(interp(frame_idx / w_frames)).to(device)
+                img = G.synthesis(ws=w.unsqueeze(0), noise_mode='const')[0]
+                imgs.append(img)
+        video_out.append_data(layout_grid(torch.stack(imgs), grid_w=grid_w, grid_h=grid_h))
+    video_out.close()
+
+
+# ----------------------------------------------------------------------------
+
+
+def parse_range(s: Union[str, List[int]]) -> List[int]:
+    """Parse a comma separated list of numbers or ranges and return a list of ints.
+
+    Example: '1,2,5-10' returns [1, 2, 5, 6, 7]
+    """
+    if isinstance(s, list): return s
+    ranges = []
+    range_re = re.compile(r'^(\d+)-(\d+)$')
+    for p in s.split(','):
+        m = range_re.match(p)
+        if m:
+            ranges.extend(range(int(m.group(1)), int(m.group(2))+1))
+        else:
+            ranges.append(int(p))
+    return ranges
+
+
+# ----------------------------------------------------------------------------
+
+
+def parse_tuple(s: Union[str, Tuple[int,int]]) -> Tuple[int, int]:
+    """Parse a 'M,N' or 'MxN' integer tuple.
+
+    Example:
+        '4x2' returns (4,2)
+        '0,1' returns (0,1)
+    """
+    if isinstance(s, tuple): return s
+    m = re.match(r'^(\d+)[x,](\d+)$', s)
+    if m:
+        return (int(m.group(1)), int(m.group(2)))
+    raise ValueError(f'cannot parse tuple {s}')
+
+
+# ----------------------------------------------------------------------------
+
+
+@click.command()
+@click.option('--network', 'network_pkl', help='Network pickle filename', required=True)
+@click.option('--seeds', type=parse_range, help='List of random seeds', required=True)
+@click.option('--shuffle-seed', type=int, help='Random seed to use for shuffling seed order', default=None)
+@click.option('--grid', type=parse_tuple, help='Grid width/height, e.g. \'4x3\' (default: 1x1)', default=(1,1))
+@click.option('--num-keyframes', type=int, help='Number of seeds to interpolate through.  If not specified, determine based on the length of the seeds array given by --seeds.', default=None)
+@click.option('--w-frames', type=int, help='Number of frames to interpolate between latents', default=120)
+@click.option('--trunc', 'truncation_psi', type=float, help='Truncation psi', default=1, show_default=True)
+@click.option('--stabilize-video', is_flag=True, help='Stabilize the video by anchoring the mapping to w_avg')
+@click.option('--output', help='Output .mp4 filename', type=str, required=True, metavar='FILE')
+def generate_images(
+    network_pkl: str,
+    seeds: List[int],
+    shuffle_seed: Optional[int],
+    truncation_psi: float,
+    grid: Tuple[int,int],
+    num_keyframes: Optional[int],
+    stabilize_video: bool,
+    w_frames: int,
+    output: str
+):
+    """Render a latent vector interpolation video.
+
+    Examples:
+
+    \b
+    # Render a 4x2 grid of interpolations for seeds 0 through 31.
+    python gen_video.py --output=lerp.mp4 --trunc=1 --seeds=0-31 --grid=4x2 \\
+        --network=https://api.ngc.nvidia.com/v2/models/nvidia/research/stylegan3/versions/1/files/stylegan3-r-afhqv2-512x512.pkl
+
+    Animation length and seed keyframes:
+
+    The animation length is either determined based on the --seeds value or explicitly
+    specified using the --num-keyframes option.
+
+    When num keyframes is specified with --num-keyframes, the output video length
+    will be 'num_keyframes*w_frames' frames.
+
+    If --num-keyframes is not specified, the number of seeds given with
+    --seeds must be divisible by grid size W*H (--grid).  In this case the
+    output video length will be '# seeds/(w*h)*w_frames' frames.
+    """
+
+    print('Loading networks from "%s"...' % network_pkl)
+    device = torch.device('cuda')
+    with dnnlib.util.open_url(network_pkl) as f:
+        G = legacy.load_network_pkl(f)['G_ema'].to(device) # type: ignore
+
+    gen_interp_video(G=G, mp4=output, bitrate='12M', grid_dims=grid, num_keyframes=num_keyframes, w_frames=w_frames,
+                     seeds=seeds, shuffle_seed=shuffle_seed, psi=truncation_psi, stabilize_video=stabilize_video)
+
+
+# ----------------------------------------------------------------------------
+
+
+if __name__ == "__main__":
+    generate_images() # pylint: disable=no-value-for-parameter
+
+
+# ----------------------------------------------------------------------------

stylegan3-fun/generate.py

@@ -0,0 +1,838 @@
+import os
+import sys
+from typing import List, Optional, Union, Tuple
+import click
+
+import dnnlib
+from torch_utils import gen_utils
+import copy
+
+import scipy
+import numpy as np
+import PIL.Image
+import torch
+
+import legacy
+from viz.renderer import Renderer
+
+os.environ['PYGAME_HIDE_SUPPORT_PROMPT'] = 'hide'
+import moviepy.editor
+
+
+# ----------------------------------------------------------------------------
+
+
+# We group the different types of generation (images, grid, video, wacky stuff) into a main function
+@click.group()
+def main():
+    pass
+
+
+# ----------------------------------------------------------------------------
+
+
+@main.command(name='images')
+@click.pass_context
+@click.option('--network', 'network_pkl', help='Network pickle filename: can be URL, local file, or the name of the model in torch_utils.gen_utils.resume_specs', required=True)
+@click.option('--device', help='Device to use for image generation; using the CPU is slower than the GPU', type=click.Choice(['cpu', 'cuda']), default='cuda', show_default=True)
+@click.option('--cfg', type=click.Choice(gen_utils.available_cfgs), help='Config of the network, used only if you want to use the pretrained models in torch_utils.gen_utils.resume_specs')
+# Synthesis options (feed a list of seeds or give the projected w to synthesize)
+@click.option('--seeds', type=gen_utils.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)', default=None, show_default=True)
+@click.option('--noise-mode', help='Noise mode', type=click.Choice(['const', 'random', 'none']), default='const', show_default=True)
+@click.option('--anchor-latent-space', '-anchor', is_flag=True, help='Anchor the latent space to w_avg to stabilize the video')
+@click.option('--projected-w', help='Projection result file; can be either .npy or .npz files', type=click.Path(exists=True, dir_okay=False), metavar='FILE')
+@click.option('--new-center', type=gen_utils.parse_new_center, help='New center for the W latent space; a seed (int) or a path to a projected dlatent (.npy/.npz)', default=None)
+# Save the output of the intermediate layers
+@click.option('--layer', 'layer_name', type=str, help='Layer name to extract; if unsure, use `--available-layers`', default=None, show_default=True)
+@click.option('--available-layers', is_flag=True, help='List the available layers in the trained model and exit')
+@click.option('--starting-channel', 'starting_channel', type=int, help='Starting channel for the layer extraction', default=0, show_default=True)
+@click.option('--grayscale', 'save_grayscale', type=bool, help='Use the first channel starting from `--starting-channel` to generate a grayscale image.', default=False, show_default=True)
+@click.option('--rgb', 'save_rgb', type=bool, help='Use 3 consecutive channels (if they exist) to generate a RGB image, starting from `--starting-channel`.', default=False, show_default=True)
+@click.option('--rgba', 'save_rgba', type=bool, help='Use 4 consecutive channels (if they exist) to generate a RGBA image, starting from `--starting-channel`.', default=False, show_default=True)
+@click.option('--img-scale-db', 'img_scale_db', type=click.FloatRange(min=-40, max=40), help='Scale the image pixel values, akin to "exposure" (lower, the image is grayer/, higher the more white/burnt regions)', default=0, show_default=True)
+@click.option('--img-normalize', 'img_normalize', type=bool, help='Normalize images of the selected layer and channel', default=False, show_default=True)
+# Grid options
+@click.option('--save-grid', is_flag=True, help='Use flag to save image grid')
+@click.option('--grid-width', '-gw', type=click.IntRange(min=1), help='Grid width (number of columns)', default=None)
+@click.option('--grid-height', '-gh', type=click.IntRange(min=1), help='Grid height (number of rows)', default=None)
+# Extra parameters for saving the results
+@click.option('--save-dlatents', is_flag=True, help='Use flag to save individual dlatents (W) for each individual resulting image')
+@click.option('--outdir', type=click.Path(file_okay=False), help='Directory path to save the results', default=os.path.join(os.getcwd(), 'out', 'images'), show_default=True, metavar='DIR')
+@click.option('--description', '-desc', type=str, help='Description name for the directory path to save results', default='generate-images', show_default=True)
+def generate_images(
+        ctx: click.Context,
+        network_pkl: str,
+        device: Optional[str],
+        cfg: Optional[str],
+        seeds: Optional[List[int]],
+        truncation_psi: Optional[float],
+        class_idx: Optional[int],
+        noise_mode: Optional[str],
+        anchor_latent_space: Optional[bool],
+        projected_w: Optional[Union[str, os.PathLike]],
+        new_center: Tuple[str, Union[int, np.ndarray]],
+        layer_name: Optional[str],
+        available_layers: Optional[bool],
+        starting_channel: Optional[int],
+        save_grayscale: Optional[bool],
+        save_rgb: Optional[bool],
+        save_rgba: Optional[bool],
+        img_scale_db: Optional[float],
+        img_normalize: Optional[bool],
+        save_grid: Optional[bool],
+        grid_width: int,
+        grid_height: int,
+        save_dlatents: Optional[bool],
+        outdir: Union[str, os.PathLike],
+        description: str,
+):
+    """Generate images using pretrained network pickle.
+
+    Examples:
+
+    \b
+    # Generate curated MetFaces images without truncation (Fig.10 left)
+    python generate.py images --trunc=1 --seeds=85,265,297,849 \\
+        --network=https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/metfaces.pkl
+
+    \b
+    # Generate uncurated MetFaces images with truncation (Fig.12 upper left)
+    python generate.py images --trunc=0.7 --seeds=600-605 \\
+        --network=https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/metfaces.pkl
+
+    \b
+    # Generate class conditional CIFAR-10 images (Fig.17 left, Car)
+    python generate.py images --seeds=0-35 --class=1 \\
+        --network=https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/cifar10.pkl
+
+    \b
+    # Render an image from projected W
+    python generate.py images --projected_w=projected_w.npz \\
+        --network=https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/metfaces.pkl
+
+    \b
+    Generate class conditional StyleGAN2 WikiArt images, save each individually, and save all of them as a grid
+    python generate.py images --cfg=stylegan2 --network=wikiart1024-C --class=155 \\
+        --trunc=0.7 --seeds=10-50 --save-grid
+    """
+    # Sanity check
+    if len(seeds) < 1:
+        ctx.fail('Use `--seeds` to specify at least one seed.')
+
+    device = torch.device('cuda') if torch.cuda.is_available() and device == 'cuda' else torch.device('cpu')
+
+    # Load the network
+    G = gen_utils.load_network('G_ema', network_pkl, cfg, device)
+
+    if available_layers:
+        click.secho(f'Printing available layers (name, channels and size) for "{network_pkl}"...', fg='blue')
+        _ = Renderer().render(G=G, available_layers=available_layers)
+        sys.exit(1)
+
+    # Setup for using CPU
+    if device.type == 'cpu':
+        gen_utils.use_cpu(G)
+
+    # Stabilize/anchor the latent space
+    if anchor_latent_space:
+        gen_utils.anchor_latent_space(G)
+
+    description = 'generate-images' if len(description) == 0 else description
+    description = f'{description}-{layer_name}_layer' if layer_name is not None else description
+    # Create the run dir with the given name description
+    run_dir = gen_utils.make_run_dir(outdir, description)
+
+    # Synthesize the result of a W projection.
+    if projected_w is not None:
+        if seeds is not None:
+            print('warn: --seeds is ignored when using --projected-w')
+        print(f'Generating images from projected W "{projected_w}"')
+        ws, ext = gen_utils.get_latent_from_file(projected_w, return_ext=True)
+        ws = torch.tensor(ws, device=device)
+        assert ws.shape[1:] == (G.num_ws, G.w_dim)
+        n_digits = int(np.log10(len(ws))) + 1  # number of digits for naming the images
+        if ext == '.npy':
+            img = gen_utils.w_to_img(G, ws, noise_mode)[0]
+            PIL.Image.fromarray(img, gen_utils.channels_dict[G.synthesis.img_channels]).save(f'{run_dir}/proj.png')
+        else:
+            for idx, w in enumerate(ws):
+                img = gen_utils.w_to_img(G, w, noise_mode)[0]
+                PIL.Image.fromarray(img,
+                                    gen_utils.channels_dict[G.synthesis.img_channels]).save(f'{run_dir}/proj{idx:0{n_digits}d}.png')
+        return
+
+    # Labels.
+    class_idx = gen_utils.parse_class(G, class_idx, ctx)
+    label = torch.zeros([1, G.c_dim], device=device)
+    if G.c_dim != 0:
+        label[:, class_idx] = 1
+    else:
+        if class_idx is not None:
+            print('warn: --class=lbl ignored when running on an unconditional network')
+
+    if seeds is None:
+        ctx.fail('--seeds option is required when not using --projected-w')
+
+    # Recenter the latent space, if specified
+    if new_center is None:
+        w_avg = G.mapping.w_avg
+    else:
+        new_center, new_center_value = new_center
+        # We get the new center using the int (a seed) or recovered dlatent (an np.ndarray)
+        if isinstance(new_center_value, int):
+            w_avg = gen_utils.get_w_from_seed(G, device, new_center_value,
+                                              truncation_psi=1.0)  # We want the pure dlatent
+        elif isinstance(new_center_value, np.ndarray):
+            w_avg = torch.from_numpy(new_center_value).to(device)
+        else:
+            ctx.fail('Error: New center has strange format! Only an int (seed) or a file (.npy/.npz) are accepted!')
+
+    # Generate images.
+    images = []
+    for seed_idx, seed in enumerate(seeds):
+        print('Generating image for seed %d (%d/%d) ...' % (seed, seed_idx, len(seeds)))
+        dlatent = gen_utils.get_w_from_seed(G, device, seed, truncation_psi=1.0)
+        # Do truncation trick with center (new or global)
+        w = w_avg + (dlatent - w_avg) * truncation_psi
+
+        # TODO: this is starting to look like an auxiliary function!
+        # Save the intermediate layer output.
+        if layer_name is not None:
+            # Sanity check (meh, could be done better)
+            submodule_names = {name: mod for name, mod in G.synthesis.named_modules()}
+            assert layer_name in submodule_names, f'Layer "{layer_name}" not found in the network! Available layers: {", ".join(submodule_names)}'
+            assert True in (save_grayscale, save_rgb, save_rgba), 'You must select to save the image in at least one of the three possible formats! (L, RGB, RGBA)'
+
+            sel_channels = 3 if save_rgb else (1 if save_grayscale else 4)
+            res = Renderer().render(G=G, layer_name=layer_name, dlatent=w, sel_channels=sel_channels,
+                                    base_channel=starting_channel, img_scale_db=img_scale_db, img_normalize=img_normalize)
+            img = res.image
+        else:
+            img = gen_utils.w_to_img(G, w, noise_mode)[0]
+
+        if save_grid:
+            images.append(img)
+
+        # Get the image format, whether user-specified or the one from the model
+        try:
+            img_format = gen_utils.channels_dict[sel_channels]
+        except NameError:
+            img_format = gen_utils.channels_dict[G.synthesis.img_channels]
+
+        # Save image, avoiding grayscale errors in PIL
+        PIL.Image.fromarray(img[:, :, 0] if img.shape[-1] == 1 else img,
+                            img_format).save(os.path.join(run_dir, f'seed{seed}.png'))
+        if save_dlatents:
+            np.save(os.path.join(run_dir, f'seed{seed}.npy'), w.unsqueeze(0).cpu().numpy())
+
+    if save_grid:
+        print('Saving image grid...')
+        images = np.array(images)
+
+        # We let the function infer the shape of the grid
+        if (grid_width, grid_height) == (None, None):
+            grid = gen_utils.create_image_grid(images)
+        # The user tells the specific shape of the grid, but one value may be None
+        else:
+            grid = gen_utils.create_image_grid(images, (grid_width, grid_height))
+
+        grid = grid[:, :, 0] if grid.shape[-1] == 1 else grid
+        PIL.Image.fromarray(grid, img_format).save(os.path.join(run_dir, 'grid.png'))
+
+    # Save the configuration used
+    ctx.obj = {
+        'network_pkl': network_pkl,
+        'device': device.type,
+        'config': cfg,
+        'synthesis': {
+            'seeds': seeds,
+            'truncation_psi': truncation_psi,
+            'class_idx': class_idx,
+            'noise_mode': noise_mode,
+            'anchor_latent_space': anchor_latent_space,
+            'projected_w': projected_w,
+            'new_center': new_center
+        },
+        'intermediate_representations': {
+            'layer': layer_name,
+            'starting_channel': starting_channel,
+            'grayscale': save_grayscale,
+            'rgb': save_rgb,
+            'rgba': save_rgba,
+            'img_scale_db': img_scale_db,
+            'img_normalize': img_normalize
+        },
+        'grid_options': {
+            'save_grid': save_grid,
+            'grid_width': grid_width,
+            'grid_height': grid_height,
+        },
+        'extra_parameters': {
+            'save_dlatents': save_dlatents,
+            'run_dir': run_dir,
+            'description': description,
+        }
+    }
+    gen_utils.save_config(ctx=ctx, run_dir=run_dir)
+
+
+# ----------------------------------------------------------------------------
+
+
+@main.command(name='random-video')
+@click.pass_context
+@click.option('--network', 'network_pkl', help='Network pickle filename', required=True)
+@click.option('--cfg', type=click.Choice(gen_utils.available_cfgs), help='Config of the network, used only if you want to use the pretrained models in torch_utils.gen_utils.resume_specs')
+# Synthesis options
+@click.option('--seeds', type=gen_utils.num_range, help='List of random seeds', required=True)
+@click.option('--trunc', 'truncation_psi', type=float, help='Truncation psi', default=1, show_default=True)
+@click.option('--new-center', type=gen_utils.parse_new_center, help='New center for the W latent space; a seed (int) or a path to a projected dlatent (.npy/.npz)', default=None)
+@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('--anchor-latent-space', '-anchor', is_flag=True, help='Anchor the latent space to w_avg to stabilize the video')
+# Save the output of the intermediate layers
+@click.option('--layer', 'layer_name', type=str, help='Layer name to extract; if unsure, use `--available-layers`', default=None, show_default=True)
+@click.option('--available-layers', is_flag=True, help='List the available layers in the trained model and exit')
+@click.option('--starting-channel', 'starting_channel', type=int, help='Starting channel for the layer extraction', default=0, show_default=True)
+@click.option('--grayscale', 'save_grayscale', type=bool, help='Use the first channel starting from `--starting-channel` to generate a grayscale image.', default=False, show_default=True)
+@click.option('--rgb', 'save_rgb', type=bool, help='Use 3 consecutive channels (if they exist) to generate a RGB image, starting from `--starting-channel`.', default=False, show_default=True)
+@click.option('--img-scale-db', 'img_scale_db', type=click.FloatRange(min=-40, max=40), help='Scale the image pixel values, akin to "exposure" (lower, the image is grayer/, higher the more white/burnt regions)', default=0, show_default=True)
+@click.option('--img-normalize', 'img_normalize', type=bool, help='Normalize images of the selected layer and channel', default=False, show_default=True)
+# Video options
+@click.option('--grid-width', '-gw', type=click.IntRange(min=1), help='Video grid width / number of columns', default=None, show_default=True)
+@click.option('--grid-height', '-gh', type=click.IntRange(min=1), help='Video grid height / number of rows', default=None, show_default=True)
+@click.option('--slowdown', type=gen_utils.parse_slowdown, help='Slow down the video by this amount; will be approximated to the nearest power of 2', default='1', show_default=True)
+@click.option('--duration-sec', '-sec', type=float, help='Duration length of the video', default=30.0, show_default=True)
+@click.option('--fps', type=click.IntRange(min=1), help='Video FPS.', default=30, show_default=True)
+@click.option('--compress', is_flag=True, help='Add flag to compress the final mp4 file with ffmpeg-python (same resolution, lower file size)')
+# Extra parameters for saving the results
+@click.option('--outdir', type=click.Path(file_okay=False), help='Directory path to save the results', default=os.path.join(os.getcwd(), 'out', 'video'), show_default=True, metavar='DIR')
+@click.option('--description', '-desc', type=str, help='Description name for the directory path to save results')
+def random_interpolation_video(
+        ctx: click.Context,
+        network_pkl: Union[str, os.PathLike],
+        cfg: Optional[str],
+        seeds: List[int],
+        truncation_psi: Optional[float],
+        new_center: Tuple[str, Union[int, np.ndarray]],
+        class_idx: Optional[int],
+        noise_mode: Optional[str],
+        anchor_latent_space: Optional[bool],
+        layer_name: Optional[str],
+        available_layers: Optional[bool],
+        starting_channel: Optional[int],
+        save_grayscale: Optional[bool],
+        save_rgb: Optional[bool],
+        img_scale_db: Optional[float],
+        img_normalize: Optional[bool],
+        grid_width: int,
+        grid_height: int,
+        slowdown: Optional[int],
+        duration_sec: Optional[float],
+        fps: int,
+        outdir: Union[str, os.PathLike],
+        description: str,
+        compress: bool,
+        smoothing_sec: Optional[float] = 3.0  # for Gaussian blur; won't be a command-line parameter, change at own risk
+):
+    """
+    Generate a random interpolation video using a pretrained network.
+
+    Examples:
+
+    \b
+    # Generate a 30-second long, untruncated MetFaces video at 30 FPS (3 rows and 2 columns; horizontal):
+    python generate.py random-video --seeds=0-5 \\
+        --network=https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/metfaces.pkl
+
+    \b
+    # Generate a 60-second long, truncated 1x2 MetFaces video at 60 FPS (2 rows and 1 column; vertical):
+    python generate.py random-video --trunc=0.7 --seeds=10,20 --grid-width=1 --grid-height=2 \\
+        --fps=60 -sec=60 --network=https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/metfaces.pkl
+
+    """
+    # Sanity check
+    if len(seeds) < 1:
+        ctx.fail('Use `--seeds` to specify at least one seed.')
+
+    device = torch.device('cuda')
+
+    # Load the network
+    G = gen_utils.load_network('G_ema', network_pkl, cfg, device)
+
+    # Print the available layers in the model
+    if available_layers:
+        click.secho(f'Printing available layers (name, channels and size) for "{network_pkl}"...', fg='blue')
+        _ = Renderer().render(G=G, available_layers=available_layers)
+        sys.exit(1)
+
+    # Sadly, render can only generate one image at a time, so for now we'll just use the first seed
+    if layer_name is not None and len(seeds) > 1:
+        print(f'Note: Only one seed is supported for layer extraction, using seed "{seeds[0]}"...')
+        seeds = seeds[:1]
+
+    # Stabilize/anchor the latent space
+    if anchor_latent_space:
+        gen_utils.anchor_latent_space(G)
+
+    # Create the run dir with the given name description; add slowdown if different from the default (1)
+    desc = 'random-video'
+    desc = f'random-video-{description}' if description is not None else desc
+    desc = f'{desc}-{slowdown}xslowdown' if slowdown != 1 else desc
+    desc = f'{desc}-{layer_name}_layer' if layer_name is not None else desc
+    run_dir = gen_utils.make_run_dir(outdir, desc)
+
+    # Number of frames in the video and its total duration in seconds
+    num_frames = int(np.rint(duration_sec * fps))
+    total_duration = duration_sec * slowdown
+
+    print('Generating latent vectors...')
+    # TODO: let another helper function handle each case, we will use it for the grid
+    # If there's more than one seed provided and the shape isn't specified by the user
+    if (grid_width is None and grid_height is None) and len(seeds) >= 1:
+        # TODO: this can be done by another function
+        # Number of images in the grid video according to the seeds provided
+        num_seeds = len(seeds)
+        # Get the grid width and height according to num, giving priority to the number of columns
+        grid_width = max(int(np.ceil(np.sqrt(num_seeds))), 1)
+        grid_height = max((num_seeds - 1) // grid_width + 1, 1)
+        grid_size = (grid_width, grid_height)
+        shape = [num_frames, G.z_dim]  # This is per seed
+        # Get the z latents
+        all_latents = np.stack([np.random.RandomState(seed).randn(*shape).astype(np.float32) for seed in seeds], axis=1)
+
+    # If only one seed is provided, but the user specifies the grid shape:
+    elif None not in (grid_width, grid_height) and len(seeds) == 1:
+        grid_size = (grid_width, grid_height)
+        shape = [num_frames, np.prod(grid_size), G.z_dim]
+        # Since we have one seed, we use it to generate all latents
+        all_latents = np.random.RandomState(*seeds).randn(*shape).astype(np.float32)
+
+    # If one or more seeds are provided, and the user also specifies the grid shape:
+    elif None not in (grid_width, grid_height) and len(seeds) >= 1:
+        # Case is similar to the first one
+        num_seeds = len(seeds)
+        grid_size = (grid_width, grid_height)
+        available_slots = np.prod(grid_size)
+        if available_slots < num_seeds:
+            diff = num_seeds - available_slots
+            click.secho(f'More seeds were provided ({num_seeds}) than available spaces in the grid ({available_slots})',
+                        fg='red')
+            click.secho(f'Removing the last {diff} seeds: {seeds[-diff:]}', fg='blue')
+            seeds = seeds[:available_slots]
+        shape = [num_frames, G.z_dim]
+        all_latents = np.stack([np.random.RandomState(seed).randn(*shape).astype(np.float32) for seed in seeds], axis=1)
+
+    else:
+        ctx.fail('Error: wrong combination of arguments! Please provide either a list of seeds, one seed and the grid '
+                 'width and height, or more than one seed and the grid width and height')
+
+    # Let's smooth out the random latents so that now they form a loop (and are correctly generated in a 512-dim space)
+    all_latents = scipy.ndimage.gaussian_filter(all_latents, sigma=[smoothing_sec * fps, 0, 0], mode='wrap')
+    all_latents /= np.sqrt(np.mean(np.square(all_latents)))
+
+    # Name of the video
+    mp4_name = f'{grid_width}x{grid_height}-slerp-{slowdown}xslowdown'
+
+    # Labels.
+    class_idx = gen_utils.parse_class(G, class_idx, ctx)
+    label = torch.zeros([1, G.c_dim], device=device)
+    if G.c_dim != 0:
+        label[:, class_idx] = 1
+    else:
+        if class_idx is not None:
+            print('warn: --class=lbl ignored when running on an unconditional network')
+
+    # Let's slowdown the video, if so desired
+    while slowdown > 1:
+        all_latents, duration_sec, num_frames = gen_utils.double_slowdown(latents=all_latents,
+                                                                          duration=duration_sec,
+                                                                          frames=num_frames)
+        slowdown //= 2
+
+    if new_center is None:
+        w_avg = G.mapping.w_avg
+    else:
+        new_center, new_center_value = new_center
+        # We get the new center using the int (a seed) or recovered dlatent (an np.ndarray)
+        if isinstance(new_center_value, int):
+            w_avg = gen_utils.get_w_from_seed(G, device, new_center_value,
+                                              truncation_psi=1.0)  # We want the pure dlatent
+        elif isinstance(new_center_value, np.ndarray):
+            w_avg = torch.from_numpy(new_center_value).to(device)
+        else:
+            ctx.fail('Error: New center has strange format! Only an int (seed) or a file (.npy/.npz) are accepted!')
+
+    # Auxiliary function for moviepy
+    def make_frame(t):
+        frame_idx = int(np.clip(np.round(t * fps), 0, num_frames - 1))
+        latents = torch.from_numpy(all_latents[frame_idx]).to(device)
+        # Do the truncation trick (with the global centroid or the new center provided by the user)
+        w = G.mapping(latents, None)
+        w = w_avg + (w - w_avg) * truncation_psi
+
+        # Get the images
+
+        # Save the intermediate layer output.
+        if layer_name is not None:
+            # Sanity check (again, could be done better)
+            submodule_names = {name: mod for name, mod in G.synthesis.named_modules()}
+            assert layer_name in submodule_names, f'Layer "{layer_name}" not found in the network! Available layers: {", ".join(submodule_names)}'
+            assert True in (save_grayscale, save_rgb), 'You must select to save the video in at least one of the two possible formats! (L, RGB)'
+
+            sel_channels = 3 if save_rgb else 1
+            res = Renderer().render(G=G, layer_name=layer_name, dlatent=w, sel_channels=sel_channels,
+                                    base_channel=starting_channel, img_scale_db=img_scale_db, img_normalize=img_normalize)
+            images = res.image
+            images = np.expand_dims(np.array(images), axis=0)
+        else:
+            images = gen_utils.w_to_img(G, w, noise_mode)  # Remember, it can only be a single image
+            # RGBA -> RGB, if necessary
+            images = images[:, :, :, :3]
+
+        # Generate the grid for this timestamp
+        grid = gen_utils.create_image_grid(images, grid_size)
+        # moviepy.editor.VideoClip expects 3 channels
+        if grid.shape[2] == 1:
+            grid = grid.repeat(3, 2)
+        return grid
+
+    # Generate video using the respective make_frame function
+    videoclip = moviepy.editor.VideoClip(make_frame, duration=duration_sec)
+    videoclip.set_duration(total_duration)
+
+    mp4_name = f'{mp4_name}_{layer_name}' if layer_name is not None else mp4_name
+
+    # Change the video parameters (codec, bitrate) if you so desire
+    final_video = os.path.join(run_dir, f'{mp4_name}.mp4')
+    videoclip.write_videofile(final_video, fps=fps, codec='libx264', bitrate='16M')
+
+    # Save the configuration used
+    new_center = 'w_avg' if new_center is None else new_center
+    ctx.obj = {
+        'network_pkl': network_pkl,
+        'config': cfg,
+        'synthesis_options': {
+            'seeds': seeds,
+            'truncation_psi': truncation_psi,
+            'new_center': new_center,
+            'class_idx': class_idx,
+            'noise_mode': noise_mode,
+            'anchor_latent_space': anchor_latent_space
+        },
+        'intermediate_representations': {
+            'layer': layer_name,
+            'starting_channel': starting_channel,
+            'grayscale': save_grayscale,
+            'rgb': save_rgb,
+            'img_scale_db': img_scale_db,
+            'img_normalize': img_normalize
+        },
+        'video_options': {
+            'grid_width': grid_width,
+            'grid_height': grid_height,
+            'slowdown': slowdown,
+            'duration_sec': duration_sec,
+            'video_fps': fps,
+            'compress': compress,
+            'smoothing_sec': smoothing_sec
+        },
+        'extra_parameters': {
+            'run_dir': run_dir,
+            'description': desc
+        }
+    }
+    gen_utils.save_config(ctx=ctx, run_dir=run_dir)
+
+    # Compress the video (lower file size, same resolution)
+    if compress:
+        gen_utils.compress_video(original_video=final_video, original_video_name=mp4_name, outdir=run_dir, ctx=ctx)
+
+
+# ----------------------------------------------------------------------------
+
+
+@main.command('circular-video')
+@click.pass_context
+@click.option('--network', 'network_pkl', help='Network pickle filename', required=True)
+@click.option('--cfg', type=click.Choice(gen_utils.available_cfgs), help='Config of the network, used only if you want to use the pretrained models in torch_utils.gen_utils.resume_specs')
+# Synthesis options
+@click.option('--seed', type=int, help='Random seed', required=True)
+@click.option('--trunc', 'truncation_psi', type=float, help='Truncation psi', default=1, show_default=True)
+@click.option('--trunc-start', 'truncation_psi_start', type=float, help='Initial value of pulsating truncation psi', default=None, show_default=True)
+@click.option('--trunc-end', 'truncation_psi_end', type=float, help='Maximum/minimum value of pulsating truncation psi', default=None, show_default=True)
+@click.option('--global-pulse', 'global_pulsation_trick', is_flag=True, help='If set, the truncation psi will pulsate globally (on all grid cells)')
+@click.option('--wave-pulse', 'wave_pulsation_trick', is_flag=True, help='If set, the truncation psi will pulsate in a wave-like fashion from the upper left to the lower right in the grid')
+@click.option('--frequency', 'pulsation_frequency', type=int, help='Frequency of the pulsation', default=1, show_default=True)
+@click.option('--new-center', type=str, help='New center for the W latent space; a seed (int) or a path to a projected dlatent (.npy/.npz)', default=None)
+@click.option('--new-w-avg', 'new_w_avg', type=gen_utils.parse_new_center, help='Path to a new "global" w_avg (seed or .npy/.npz file) to be used in the truncation trick', default=None)
+@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('--anchor-latent-space', '-anchor', is_flag=True, help='Anchor the latent space to w_avg to stabilize the video')
+@click.option('--flesh', 'aydao_flesh_digression', is_flag=True, help='If set, we will slowly modify the constant input to the network (based on @aydao\'s work')
+# Video options
+@click.option('--grid-width', '-gw', type=click.IntRange(min=1), help='Video grid width / number of columns', required=True)
+@click.option('--grid-height', '-gh', type=click.IntRange(min=1), help='Video grid height / number of rows', required=True)
+@click.option('--duration-sec', '-sec', type=float, help='Duration length of the video', default=10.0, show_default=True)
+@click.option('--fps', type=click.IntRange(min=1), help='Video FPS.', default=30, show_default=True)
+@click.option('--compress', is_flag=True, help='Add flag to compress the final mp4 file with ffmpeg-python (same resolution, lower file size)')
+# Extra parameters for saving the results
+@click.option('--outdir', type=click.Path(file_okay=False), help='Directory path to save the results', default=os.path.join(os.getcwd(), 'out', 'video'), show_default=True, metavar='DIR')
+@click.option('--description', '-desc', type=str, help='Description name for the directory path to save results')
+def circular_video(
+        ctx: click.Context,
+        network_pkl: Union[str, os.PathLike],
+        cfg: Optional[str],
+        seed: int,
+        truncation_psi: Optional[float],
+        truncation_psi_start: Optional[float],
+        truncation_psi_end: Optional[float],
+        global_pulsation_trick: Optional[bool],
+        wave_pulsation_trick: Optional[bool],
+        pulsation_frequency: Optional[int],
+        new_center: Tuple[str, Union[int, np.ndarray]],
+        new_w_avg: Optional[Union[str, os.PathLike]],
+        class_idx: Optional[int],
+        noise_mode: Optional[str],
+        anchor_latent_space: Optional[bool],
+        aydao_flesh_digression: Optional[bool],
+        grid_width: int,
+        grid_height: int,
+        duration_sec: float,
+        fps: int,
+        compress: Optional[bool],
+        outdir: Union[str, os.PathLike],
+        description: str
+):
+    """
+    Generate a circular interpolation video in two random axes of Z, given a seed
+    """
+
+    device = torch.device('cuda')
+
+    # Load the network
+    G = gen_utils.load_network('G_ema', network_pkl, cfg, device)
+
+    # Get the constant input
+    if aydao_flesh_digression:
+        if hasattr(G.synthesis, 'b4'):
+            model_type = 'stylegan2'
+            const_input = copy.deepcopy(G.synthesis.b4.const).cpu().numpy()
+        elif hasattr(G.synthesis, 'input'):
+            model_type = 'stylegan3'
+            input_frequencies = copy.deepcopy(G.synthesis.input.freqs).cpu().numpy()
+            input_phases = copy.deepcopy(G.synthesis.input.phases).cpu().numpy()
+        else:
+            ctx.fail('Error: This option is only available for StyleGAN2 and StyleGAN3 models!')
+
+    # Get the labels, if the model is conditional
+    class_idx = gen_utils.parse_class(G, class_idx, ctx)
+    label = torch.zeros([1, G.c_dim], device=device)
+    if G.c_dim != 0:
+        label[:, class_idx] = 1
+    else:
+        if class_idx is not None:
+            print('warn: --class=lbl ignored when running on an unconditional network')
+
+    # Get center of the latent space (global or user-indicated)
+    if new_center is None:
+        w_avg = G.mapping.w_avg
+        w_avg = w_avg.view(1, 1, -1)  # [w_dim] => [1, 1, w_dim]
+    else:
+        # It's an int, so use as a seed
+        if new_center.isdigit():
+            w_avg = gen_utils.get_w_from_seed(G, device, int(new_center), truncation_psi=1.0).to(device)
+        # It's a file, so load it
+        elif os.path.isfile(new_center):
+            w_avg = gen_utils.get_latent_from_file(new_center, return_ext=False)
+            w_avg = torch.from_numpy(w_avg).to(device)
+        # It's a directory, so get all latents inside it (including subdirectories, so be careful)
+        elif os.path.isdir(new_center):
+            w_avg = gen_utils.parse_all_projected_dlatents(new_center)
+            w_avg = torch.tensor(w_avg).squeeze(1).to(device)
+        else:
+            message = 'Only seeds (int) or paths to latent files (.npy/.npz) or directories containing these are allowed for "--new-center"'
+            raise ctx.fail(message)
+
+        # Some sanity checks
+        num_centers = len(w_avg)
+        if num_centers == 0:
+            raise ctx.fail('No centers were found! If files, makes sure they are .npy or .npz files.')
+        # Just one is provided, so this will be a sort of 'global' center
+        elif num_centers == 1:
+            print(f'Using only one center (if more than one is desired, provide a directory with all of them)')
+        elif num_centers != grid_height * grid_width:
+            message = f"Number of centers ({num_centers}) doesn't match the grid size ({grid_height}x{grid_width})"
+            raise ctx.fail(message)
+
+    print('Using wave pulsation trick' if wave_pulsation_trick else 'Using global pulsation trick' if global_pulsation_trick else 'Using standard truncation trick...')
+    # Stabilize/anchor the latent space
+    if anchor_latent_space:
+        gen_utils.anchor_latent_space(G)
+
+    # Create the run dir with the given name description; add slowdown if different from the default (1)
+    desc = 'circular-video'
+    desc = f'circular-video-{description}' if description is not None else desc
+    desc = f'{desc}-aydao-flesh-digression' if aydao_flesh_digression else desc
+    run_dir = gen_utils.make_run_dir(outdir, desc)
+
+    # Calculate the total number of frames in the video
+    num_frames = int(np.rint(duration_sec * fps))
+
+    grid_size = (grid_width, grid_height)
+    # Get the latents with the random state
+    random_state = np.random.RandomState(seed)
+    # Choose two random dims on which to plot the circles (from 0 to G.z_dim-1),
+    # one pair for each element of the grid (2*grid_width*grid_height in total)
+    try:
+        z1, z2 = np.split(random_state.choice(G.z_dim, 2 * np.prod(grid_size), replace=False), 2)
+    except ValueError:
+        # Extreme case: G.z_dim < 2 * grid_width * grid_height (low G.z_dim most likely)
+        z1, z2 = np.split(random_state.choice(G.z_dim, 2 * np.prod(grid_size), replace=True), 2)
+
+    # We partition the circle in equal strides w.r.t. num_frames
+    get_angles = lambda num_frames: np.linspace(0, 2*np.pi, num_frames)
+    angles = get_angles(num_frames=num_frames)
+
+    # Basic Polar to Cartesian transformation
+    polar_to_cartesian = lambda radius, theta: (radius * np.cos(theta), radius * np.sin(theta))
+    # Using a fixed radius (this value is irrelevant), we generate the circles in each chosen grid
+    Z1, Z2 = polar_to_cartesian(radius=5.0, theta=angles)
+
+    # Our latents will be comprising mostly of zeros
+    all_latents = np.zeros([num_frames, np.prod(grid_size), G.z_dim]).astype(np.float32)
+    # Obtain all the frames belonging to the specific box in the grid,
+    # replacing the zero values with the circle perimeter values
+    for box in range(np.prod(grid_size)):
+        box_frames = all_latents[:, box]
+        box_frames[:, [z1[box], z2[box]]] = np.vstack((Z1, Z2)).T
+
+    if aydao_flesh_digression:
+        # We will modify the constant input to the network (for --cfg=stylegan2)
+        if model_type == 'stylegan2':
+            const_input_interpolation = np.random.randn(num_frames, *const_input.shape).astype(np.float32) / 4  # [num_frames, G.w_dim, 4, 4] ;  "/ 4" is arbitrary
+            const_input_interpolation = scipy.ndimage.gaussian_filter(const_input_interpolation, sigma=[fps, 0, 0, 0], mode='wrap')
+            const_input_interpolation /= np.sqrt(np.mean(np.square(const_input_interpolation)))
+        elif model_type == 'stylegan3':
+            const_freq_interpolation = np.random.randn(num_frames, *input_frequencies.shape).astype(np.float32) / 32  # [num_frames, G.w_dim, 2]
+            const_freq_interpolation = scipy.ndimage.gaussian_filter(const_freq_interpolation, sigma=[5.0*fps, 0, 0], mode='wrap')
+            const_freq_interpolation /= np.sqrt(np.mean(np.square(const_freq_interpolation)))
+
+            const_phase_interpolation = np.random.randn(num_frames, *input_phases.shape).astype(np.float32) / 8  # [num_frames, G.w_dim, 2]
+            const_phase_interpolation = scipy.ndimage.gaussian_filter(const_phase_interpolation, sigma=[5.0*fps, 0], mode='wrap')
+            const_phase_interpolation /= np.sqrt(np.mean(np.square(const_phase_interpolation)))
+
+    # Convert to torch tensor
+    if new_w_avg is not None:
+        print("Moving all the latent space towards the new center...")
+        _, new_w_avg = new_w_avg
+        # We get the new center using the int (a seed) or recovered dlatent (an np.ndarray)
+        if isinstance(new_w_avg, int):
+            new_w_avg = gen_utils.get_w_from_seed(G, device, new_w_avg,
+                                                  truncation_psi=1.0)  # We want the pure dlatent
+        elif isinstance(new_w_avg, np.ndarray):
+            new_w_avg = torch.from_numpy(new_w_avg).to(device)  # [1, num_ws, w_dim]
+        else:
+            ctx.fail('Error: New center has strange format! Only an int (seed) or a file (.npy/.npz) are accepted!')
+
+    # Auxiliary function for moviepy
+    def make_frame(t):
+        frame_idx = int(np.clip(np.round(t * fps), 0, num_frames - 1))
+        latents = torch.from_numpy(all_latents[frame_idx]).to(device)
+        # Get the images with the respective label
+        dlatents = gen_utils.z_to_dlatent(G, latents, label, truncation_psi=1.0)  # Get the pure dlatent
+        # Do truncation trick
+        # For the truncation trick (supersedes any value chosen for truncation_psi)
+        if None not in (truncation_psi_start, truncation_psi_end):
+            # For both, truncation psi will have the general form of a sinusoid: psi = (cos(t) + alpha) / beta
+            if global_pulsation_trick:
+                tr = gen_utils.global_pulsate_psi(psi_start=truncation_psi_start,
+                                                  psi_end=truncation_psi_end,
+                                                  n_steps=num_frames)
+            elif wave_pulsation_trick:
+                tr = gen_utils.wave_pulse_truncation_psi(psi_start=truncation_psi_start,
+                                                         psi_end=truncation_psi_end,
+                                                         n_steps=num_frames,
+                                                         grid_shape=grid_size,
+                                                         frequency=pulsation_frequency,
+                                                         time=frame_idx)
+        # Define how to use the truncation psi
+        if global_pulsation_trick:
+            tr = tr[frame_idx].to(device)
+        elif wave_pulsation_trick:
+            tr = tr.to(device)
+        else:
+            # It's a float, so we can just use it
+            tr = truncation_psi
+
+        w = w_avg + (dlatents - w_avg) * tr
+        # Modify the constant input
+        if aydao_flesh_digression:
+            if model_type == 'stylegan2':
+                G.synthesis.b4.const.copy_(torch.from_numpy(const_input_interpolation[frame_idx]))
+            elif model_type == 'stylegan3':
+                pass
+                # G.synthesis.input.freqs.copy_(torch.from_numpy(const_freq_interpolation[frame_idx]))
+                # G.synthesis.input.phases.copy_(torch.from_numpy(const_phase_interpolation[frame_idx]))
+                # G.synthesis.input.phases.copy_(torch.from_numpy(
+                #     input_phases * np.cos(np.pi * frame_idx / num_frames) ** 2
+                #     ))
+        # Get the images
+        images = gen_utils.w_to_img(G, w, noise_mode, new_w_avg, tr)
+        # RGBA -> RGB
+        images = images[:, :, :, :3]
+        # Generate the grid for this timestep
+        grid = gen_utils.create_image_grid(images, grid_size)
+        # Grayscale => RGB
+        if grid.shape[2] == 1:
+            grid = grid.repeat(3, 2)
+        return grid
+
+    # Generate video using the respective make_frame function
+    videoclip = moviepy.editor.VideoClip(make_frame, duration=duration_sec)
+    videoclip.set_duration(duration_sec)
+
+    # Name of the video
+    mp4_name = f'{grid_width}x{grid_height}-circular'
+
+    # Change the video parameters (codec, bitrate) if you so desire
+    final_video = os.path.join(run_dir, f'{mp4_name}.mp4')
+    videoclip.write_videofile(final_video, fps=fps, codec='libx264', bitrate='16M')
+
+    # Save the configuration used
+    new_center = 'w_avg' if new_center is None else new_center
+    ctx.obj = {
+        'network_pkl': network_pkl,
+        'config': cfg,
+        'seed': seed,
+        'z1, z2': [[int(i), int(j)] for i, j in zip(z1, z2)],
+        'truncation_psi': truncation_psi if isinstance(truncation_psi, float) else 'pulsating',
+        'truncation_psi_start': truncation_psi_start,
+        'truncation_psi_end': truncation_psi_end,
+        'new_center': new_center,
+        'class_idx': class_idx,
+        'noise_mode': noise_mode,
+        'grid_width': grid_width,
+        'grid_height': grid_height,
+        'duration_sec': duration_sec,
+        'video_fps': fps,
+        'run_dir': run_dir,
+        'description': desc,
+        'compress': compress
+    }
+    gen_utils.save_config(ctx=ctx, run_dir=run_dir)
+
+    # Compress the video (lower file size, same resolution)
+    if compress:
+        gen_utils.compress_video(original_video=final_video, original_video_name=mp4_name, outdir=run_dir, ctx=ctx)
+
+# ----------------------------------------------------------------------------
+
+
+if __name__ == "__main__":
+    main()  # pylint: disable=no-value-for-parameter
+
+
+# ----------------------------------------------------------------------------

stylegan3-fun/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