Showing preview only (1,621K chars total). Download the full file or copy to clipboard to get everything.
Repository: XingangPan/DragGAN
Branch: main
Commit: 336f120ce126
Files: 186
Total size: 1.5 MB
Directory structure:
gitextract_5nncgy8o/
├── .gitignore
├── Dockerfile
├── LICENSE.txt
├── README.md
├── dnnlib/
│ ├── __init__.py
│ └── util.py
├── environment.yml
├── gen_images.py
├── gradio_utils/
│ ├── __init__.py
│ └── utils.py
├── gui_utils/
│ ├── __init__.py
│ ├── gl_utils.py
│ ├── glfw_window.py
│ ├── imgui_utils.py
│ ├── imgui_window.py
│ └── text_utils.py
├── legacy.py
├── requirements.txt
├── stylegan_human/
│ ├── .gitignore
│ ├── PP_HumanSeg/
│ │ ├── deploy/
│ │ │ └── infer.py
│ │ ├── export_model/
│ │ │ └── download_export_model.py
│ │ └── pretrained_model/
│ │ └── download_pretrained_model.py
│ ├── README.md
│ ├── __init__.py
│ ├── alignment.py
│ ├── bg_white.py
│ ├── dnnlib/
│ │ ├── __init__.py
│ │ ├── tflib/
│ │ │ ├── __init__.py
│ │ │ ├── autosummary.py
│ │ │ ├── custom_ops.py
│ │ │ ├── network.py
│ │ │ ├── ops/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── fused_bias_act.cu
│ │ │ │ ├── fused_bias_act.py
│ │ │ │ ├── upfirdn_2d.cu
│ │ │ │ └── upfirdn_2d.py
│ │ │ ├── optimizer.py
│ │ │ └── tfutil.py
│ │ └── util.py
│ ├── docs/
│ │ └── Dataset.md
│ ├── edit/
│ │ ├── __init__.py
│ │ ├── edit_config.py
│ │ └── edit_helper.py
│ ├── edit.py
│ ├── environment.yml
│ ├── generate.py
│ ├── insetgan.py
│ ├── interpolation.py
│ ├── latent_direction/
│ │ └── ss_statics/
│ │ ├── bottom_length_statis/
│ │ │ ├── 3/
│ │ │ │ └── statis.csv
│ │ │ ├── 4/
│ │ │ │ └── statis.csv
│ │ │ └── 5/
│ │ │ └── statis.csv
│ │ └── upper_length_statis/
│ │ └── 5/
│ │ └── statis.csv
│ ├── legacy.py
│ ├── openpose/
│ │ ├── model/
│ │ │ └── .gitkeep
│ │ └── src/
│ │ ├── __init__.py
│ │ ├── body.py
│ │ ├── model.py
│ │ └── util.py
│ ├── pti/
│ │ ├── pti_configs/
│ │ │ ├── __init__.py
│ │ │ ├── global_config.py
│ │ │ ├── hyperparameters.py
│ │ │ └── paths_config.py
│ │ ├── pti_models/
│ │ │ ├── __init__.py
│ │ │ └── e4e/
│ │ │ ├── __init__.py
│ │ │ ├── encoders/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── helpers.py
│ │ │ │ ├── model_irse.py
│ │ │ │ └── psp_encoders.py
│ │ │ ├── latent_codes_pool.py
│ │ │ ├── psp.py
│ │ │ └── stylegan2/
│ │ │ ├── __init__.py
│ │ │ ├── model.py
│ │ │ └── op/
│ │ │ ├── __init__.py
│ │ │ ├── fused_act.py
│ │ │ ├── fused_bias_act.cpp
│ │ │ ├── fused_bias_act_kernel.cu
│ │ │ ├── upfirdn2d.cpp
│ │ │ ├── upfirdn2d.py
│ │ │ └── upfirdn2d_kernel.cu
│ │ └── training/
│ │ ├── __init__.py
│ │ ├── coaches/
│ │ │ ├── __init__.py
│ │ │ ├── base_coach.py
│ │ │ ├── localitly_regulizer.py
│ │ │ ├── multi_id_coach.py
│ │ │ └── single_id_coach.py
│ │ └── projectors/
│ │ ├── __init__.py
│ │ ├── w_plus_projector.py
│ │ └── w_projector.py
│ ├── run_pti.py
│ ├── style_mixing.py
│ ├── stylemixing_video.py
│ ├── torch_utils/
│ │ ├── __init__.py
│ │ ├── custom_ops.py
│ │ ├── misc.py
│ │ ├── models.py
│ │ ├── models_face.py
│ │ ├── op_edit/
│ │ │ ├── __init__.py
│ │ │ ├── fused_act.py
│ │ │ ├── fused_bias_act.cpp
│ │ │ ├── fused_bias_act_kernel.cu
│ │ │ ├── upfirdn2d.cpp
│ │ │ ├── upfirdn2d.py
│ │ │ └── upfirdn2d_kernel.cu
│ │ ├── ops/
│ │ │ ├── __init__.py
│ │ │ ├── bias_act.cpp
│ │ │ ├── bias_act.cu
│ │ │ ├── bias_act.h
│ │ │ ├── bias_act.py
│ │ │ ├── conv2d_gradfix.py
│ │ │ ├── conv2d_resample.py
│ │ │ ├── filtered_lrelu.cpp
│ │ │ ├── filtered_lrelu.cu
│ │ │ ├── filtered_lrelu.h
│ │ │ ├── filtered_lrelu.py
│ │ │ ├── filtered_lrelu_ns.cu
│ │ │ ├── filtered_lrelu_rd.cu
│ │ │ ├── filtered_lrelu_wr.cu
│ │ │ ├── fma.py
│ │ │ ├── grid_sample_gradfix.py
│ │ │ ├── upfirdn2d.cpp
│ │ │ ├── upfirdn2d.cu
│ │ │ ├── upfirdn2d.h
│ │ │ └── upfirdn2d.py
│ │ ├── persistence.py
│ │ └── training_stats.py
│ ├── training/
│ │ ├── __init__.py
│ │ ├── augment.py
│ │ ├── dataset.py
│ │ ├── loss.py
│ │ ├── networks_stylegan2.py
│ │ ├── networks_stylegan3.py
│ │ └── training_loop.py
│ ├── training_scripts/
│ │ ├── sg2/
│ │ │ ├── train.py
│ │ │ └── training/
│ │ │ ├── dataset.py
│ │ │ └── networks.py
│ │ └── sg3/
│ │ ├── train.py
│ │ └── training/
│ │ ├── dataset.py
│ │ ├── networks_stylegan2.py
│ │ └── networks_stylegan3.py
│ └── utils/
│ ├── ImagesDataset.py
│ ├── __init__.py
│ ├── data_utils.py
│ ├── face_alignment.py
│ ├── log_utils.py
│ ├── models_utils.py
│ └── util.py
├── torch_utils/
│ ├── __init__.py
│ ├── custom_ops.py
│ ├── misc.py
│ ├── ops/
│ │ ├── __init__.py
│ │ ├── bias_act.cpp
│ │ ├── bias_act.cu
│ │ ├── bias_act.h
│ │ ├── bias_act.py
│ │ ├── conv2d_gradfix.py
│ │ ├── conv2d_resample.py
│ │ ├── filtered_lrelu.cpp
│ │ ├── filtered_lrelu.cu
│ │ ├── filtered_lrelu.h
│ │ ├── filtered_lrelu.py
│ │ ├── filtered_lrelu_ns.cu
│ │ ├── filtered_lrelu_rd.cu
│ │ ├── filtered_lrelu_wr.cu
│ │ ├── fma.py
│ │ ├── grid_sample_gradfix.py
│ │ ├── upfirdn2d.cpp
│ │ ├── upfirdn2d.cu
│ │ ├── upfirdn2d.h
│ │ └── upfirdn2d.py
│ ├── persistence.py
│ └── training_stats.py
├── training/
│ ├── __init__.py
│ ├── augment.py
│ ├── dataset.py
│ ├── loss.py
│ ├── networks_stylegan2.py
│ ├── networks_stylegan3.py
│ └── training_loop.py
├── visualizer_drag.py
├── visualizer_drag_gradio.py
└── viz/
├── __init__.py
├── capture_widget.py
├── drag_widget.py
├── latent_widget.py
├── pickle_widget.py
└── renderer.py
================================================
FILE CONTENTS
================================================
================================================
FILE: .gitignore
================================================
# Created by .ignore support plugin (hsz.mobi)
### Python template
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
*$py.class
# C extensions
*.so
# Distribution / packaging
.Python
env/
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
*.egg-info/
.installed.cfg
*.egg
# 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/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*,cover
.hypothesis/
# Translations
*.mo
*.pot
# Django stuff:
*.log
local_settings.py
# Flask stuff:
instance/
.webassets-cache
# Scrapy stuff:
.scrapy
# Sphinx documentation
docs/_build/
# PyBuilder
target/
# IPython Notebook
.ipynb_checkpoints
# pyenv
.python-version
# celery beat schedule file
celerybeat-schedule
# dotenv
.env
# virtualenv
venv/
ENV/
# Spyder project settings
.spyderproject
# Rope project settings
.ropeproject
### VirtualEnv template
# Virtualenv
# http://iamzed.com/2009/05/07/a-primer-on-virtualenv/
.Python
[Bb]in
[Ii]nclude
[Ll]ib
[Ll]ib64
[Ll]ocal
[Ss]cripts
!scripts/
pyvenv.cfg
.venv
pip-selfcheck.json
### JetBrains template
# Covers JetBrains IDEs: IntelliJ, RubyMine, PhpStorm, AppCode, PyCharm, CLion, Android Studio and Webstorm
# Reference: https://intellij-support.jetbrains.com/hc/en-us/articles/206544839
# User-specific stuff:
.idea/workspace.xml
.idea/tasks.xml
.idea/dictionaries
.idea/vcs.xml
.idea/jsLibraryMappings.xml
# Sensitive or high-churn files:
.idea/dataSources.ids
.idea/dataSources.xml
.idea/dataSources.local.xml
.idea/sqlDataSources.xml
.idea/dynamic.xml
.idea/uiDesigner.xml
# Gradle:
.idea/gradle.xml
.idea/libraries
# Mongo Explorer plugin:
.idea/mongoSettings.xml
.idea/
## File-based project format:
*.iws
## Plugin-specific files:
# IntelliJ
/out/
# mpeltonen/sbt-idea plugin
.idea_modules/
# JIRA plugin
atlassian-ide-plugin.xml
# Crashlytics plugin (for Android Studio and IntelliJ)
com_crashlytics_export_strings.xml
crashlytics.properties
crashlytics-build.properties
fabric.properties
# Mac related
.DS_Store
checkpoints
================================================
FILE: Dockerfile
================================================
FROM nvcr.io/nvidia/pytorch:23.05-py3
ENV PYTHONDONTWRITEBYTECODE 1
ENV PYTHONUNBUFFERED 1
RUN apt-get update && apt-get install -y --no-install-recommends \
make \
pkgconf \
xz-utils \
xorg-dev \
libgl1-mesa-dev \
libglu1-mesa-dev \
libxrandr-dev \
libxinerama-dev \
libxcursor-dev \
libxi-dev \
libxxf86vm-dev \
&& rm -rf /var/lib/apt/lists/*
RUN pip install --no-cache-dir --upgrade pip
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
WORKDIR /workspace
RUN (printf '#!/bin/bash\nexec \"$@\"\n' >> /entry.sh) && chmod a+x /entry.sh
ENTRYPOINT ["/entry.sh"]
================================================
FILE: LICENSE.txt
================================================
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.
=======================================================================
================================================
FILE: README.md
================================================
<p align="center">
<h1 align="center">Drag Your GAN: Interactive Point-based Manipulation on the Generative Image Manifold</h1>
<p align="center">
<a href="https://xingangpan.github.io/"><strong>Xingang Pan</strong></a>
·
<a href="https://ayushtewari.com/"><strong>Ayush Tewari</strong></a>
·
<a href="https://people.mpi-inf.mpg.de/~tleimkue/"><strong>Thomas Leimkühler</strong></a>
·
<a href="https://lingjie0206.github.io/"><strong>Lingjie Liu</strong></a>
·
<a href="https://www.meka.page/"><strong>Abhimitra Meka</strong></a>
·
<a href="http://www.mpi-inf.mpg.de/~theobalt/"><strong>Christian Theobalt</strong></a>
</p>
<h2 align="center">SIGGRAPH 2023 Conference Proceedings</h2>
<div align="center">
<img src="DragGAN.gif", width="600">
</div>
<p align="center">
<br>
<a href="https://pytorch.org/get-started/locally/"><img alt="PyTorch" src="https://img.shields.io/badge/PyTorch-ee4c2c?logo=pytorch&logoColor=white"></a>
<a href="https://twitter.com/XingangP"><img alt='Twitter' src="https://img.shields.io/twitter/follow/XingangP?label=%40XingangP"></a>
<a href="https://arxiv.org/abs/2305.10973">
<img src='https://img.shields.io/badge/Paper-PDF-green?style=for-the-badge&logo=adobeacrobatreader&logoWidth=20&logoColor=white&labelColor=66cc00&color=94DD15' alt='Paper PDF'>
</a>
<a href='https://vcai.mpi-inf.mpg.de/projects/DragGAN/'>
<img src='https://img.shields.io/badge/DragGAN-Page-orange?style=for-the-badge&logo=Google%20chrome&logoColor=white&labelColor=D35400' alt='Project Page'></a>
<a href="https://colab.research.google.com/drive/1mey-IXPwQC_qSthI5hO-LTX7QL4ivtPh?usp=sharing"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"></a>
</p>
</p>
## Web Demos
[](https://openxlab.org.cn/apps/detail/XingangPan/DragGAN)
<p align="left">
<a href="https://huggingface.co/spaces/radames/DragGan"><img alt="Huggingface" src="https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-DragGAN-orange"></a>
</p>
## Requirements
If you have CUDA graphic card, please follow the requirements of [NVlabs/stylegan3](https://github.com/NVlabs/stylegan3#requirements).
The usual installation steps involve the following commands, they should set up the correct CUDA version and all the python packages
```
conda env create -f environment.yml
conda activate stylegan3
```
Then install the additional requirements
```
pip install -r requirements.txt
```
Otherwise (for GPU acceleration on MacOS with Silicon Mac M1/M2, or just CPU) try the following:
```sh
cat environment.yml | \
grep -v -E 'nvidia|cuda' > environment-no-nvidia.yml && \
conda env create -f environment-no-nvidia.yml
conda activate stylegan3
# On MacOS
export PYTORCH_ENABLE_MPS_FALLBACK=1
```
## Run Gradio visualizer in Docker
Provided docker image is based on NGC PyTorch repository. To quickly try out visualizer in Docker, run the following:
```sh
# before you build the docker container, make sure you have cloned this repo, and downloaded the pretrained model by `python scripts/download_model.py`.
docker build . -t draggan:latest
docker run -p 7860:7860 -v "$PWD":/workspace/src -it draggan:latest bash
# (Use GPU)if you want to utilize your Nvidia gpu to accelerate in docker, please add command tag `--gpus all`, like:
# docker run --gpus all -p 7860:7860 -v "$PWD":/workspace/src -it draggan:latest bash
cd src && python visualizer_drag_gradio.py --listen
```
Now you can open a shared link from Gradio (printed in the terminal console).
Beware the Docker image takes about 25GB of disk space!
## Download pre-trained StyleGAN2 weights
To download pre-trained weights, simply run:
```
python scripts/download_model.py
```
If you want to try StyleGAN-Human and the Landscapes HQ (LHQ) dataset, please download weights from these links: [StyleGAN-Human](https://drive.google.com/file/d/1dlFEHbu-WzQWJl7nBBZYcTyo000H9hVm/view?usp=sharing), [LHQ](https://drive.google.com/file/d/16twEf0T9QINAEoMsWefoWiyhcTd-aiWc/view?usp=sharing), and put them under `./checkpoints`.
Feel free to try other pretrained StyleGAN.
## Run DragGAN GUI
To start the DragGAN GUI, simply run:
```sh
sh scripts/gui.sh
```
If you are using windows, you can run:
```
.\scripts\gui.bat
```
This GUI supports editing GAN-generated images. To edit a real image, you need to first perform GAN inversion using tools like [PTI](https://github.com/danielroich/PTI). Then load the new latent code and model weights to the GUI.
You can run DragGAN Gradio demo as well, this is universal for both windows and linux:
```sh
python visualizer_drag_gradio.py
```
## Acknowledgement
This code is developed based on [StyleGAN3](https://github.com/NVlabs/stylegan3). Part of the code is borrowed from [StyleGAN-Human](https://github.com/stylegan-human/StyleGAN-Human).
(cheers to the community as well)
## License
The code related to the DragGAN algorithm is licensed under [CC-BY-NC](https://creativecommons.org/licenses/by-nc/4.0/).
However, most of this project are available under a separate license terms: all codes used or modified from [StyleGAN3](https://github.com/NVlabs/stylegan3) is under the [Nvidia Source Code License](https://github.com/NVlabs/stylegan3/blob/main/LICENSE.txt).
Any form of use and derivative of this code must preserve the watermarking functionality showing "AI Generated".
## BibTeX
```bibtex
@inproceedings{pan2023draggan,
title={Drag Your GAN: Interactive Point-based Manipulation on the Generative Image Manifold},
author={Pan, Xingang and Tewari, Ayush, and Leimk{\"u}hler, Thomas and Liu, Lingjie and Meka, Abhimitra and Theobalt, Christian},
booktitle = {ACM SIGGRAPH 2023 Conference Proceedings},
year={2023}
}
```
================================================
FILE: dnnlib/__init__.py
================================================
# 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
================================================
FILE: dnnlib/util.py
================================================
# 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)
================================================
FILE: environment.yml
================================================
name: stylegan3
channels:
- pytorch
- nvidia
dependencies:
- python >= 3.8
- pip
- numpy>=1.25
- click>=8.0
- pillow=9.4.0
- scipy=1.11.1
- pytorch>=2.0.1
- torchvision>=0.15.2
- cudatoolkit=11.1
- requests=2.26.0
- tqdm=4.62.2
- ninja=1.10.2
- matplotlib=3.4.2
- imageio=2.9.0
- pip:
- imgui==2.0.0
- glfw==2.6.1
- gradio==3.35.2
- pyopengl==3.1.5
- imageio-ffmpeg==0.4.3
# pyspng is currently broken on MacOS (see https://github.com/nurpax/pyspng/pull/6 for instance)
- pyspng-seunglab
================================================
FILE: gen_images.py
================================================
# 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' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu')
dtype = torch.float32 if device.type == 'mps' else torch.float64
with dnnlib.util.open_url(network_pkl) as f:
G = legacy.load_network_pkl(f)['G_ema'].to(device, dtype=dtype) # type: ignore
# import pickle
# G = legacy.load_network_pkl(f)
# output = open('checkpoints/stylegan2-car-config-f-pt.pkl', 'wb')
# pickle.dump(G, output)
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, dtype=dtype)
# 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
#----------------------------------------------------------------------------
================================================
FILE: gradio_utils/__init__.py
================================================
from .utils import (ImageMask, draw_mask_on_image, draw_points_on_image,
get_latest_points_pair, get_valid_mask,
on_change_single_global_state)
__all__ = [
'draw_mask_on_image', 'draw_points_on_image',
'on_change_single_global_state', 'get_latest_points_pair',
'get_valid_mask', 'ImageMask'
]
================================================
FILE: gradio_utils/utils.py
================================================
import gradio as gr
import numpy as np
from PIL import Image, ImageDraw
class ImageMask(gr.components.Image):
"""
Sets: source="canvas", tool="sketch"
"""
is_template = True
def __init__(self, **kwargs):
super().__init__(source="upload",
tool="sketch",
interactive=False,
**kwargs)
def preprocess(self, x):
if x is None:
return x
if self.tool == "sketch" and self.source in ["upload", "webcam"
] and type(x) != dict:
decode_image = gr.processing_utils.decode_base64_to_image(x)
width, height = decode_image.size
mask = np.ones((height, width, 4), dtype=np.uint8)
mask[..., -1] = 255
mask = self.postprocess(mask)
x = {'image': x, 'mask': mask}
return super().preprocess(x)
def get_valid_mask(mask: np.ndarray):
"""Convert mask from gr.Image(0 to 255, RGBA) to binary mask.
"""
if mask.ndim == 3:
mask_pil = Image.fromarray(mask).convert('L')
mask = np.array(mask_pil)
if mask.max() == 255:
mask = mask / 255
return mask
def draw_points_on_image(image,
points,
curr_point=None,
highlight_all=True,
radius_scale=0.01):
overlay_rgba = Image.new("RGBA", image.size, 0)
overlay_draw = ImageDraw.Draw(overlay_rgba)
for point_key, point in points.items():
if ((curr_point is not None and curr_point == point_key)
or highlight_all):
p_color = (255, 0, 0)
t_color = (0, 0, 255)
else:
p_color = (255, 0, 0, 35)
t_color = (0, 0, 255, 35)
rad_draw = int(image.size[0] * radius_scale)
p_start = point.get("start_temp", point["start"])
p_target = point["target"]
if p_start is not None and p_target is not None:
p_draw = int(p_start[0]), int(p_start[1])
t_draw = int(p_target[0]), int(p_target[1])
overlay_draw.line(
(p_draw[0], p_draw[1], t_draw[0], t_draw[1]),
fill=(255, 255, 0),
width=2,
)
if p_start is not None:
p_draw = int(p_start[0]), int(p_start[1])
overlay_draw.ellipse(
(
p_draw[0] - rad_draw,
p_draw[1] - rad_draw,
p_draw[0] + rad_draw,
p_draw[1] + rad_draw,
),
fill=p_color,
)
if curr_point is not None and curr_point == point_key:
# overlay_draw.text(p_draw, "p", font=font, align="center", fill=(0, 0, 0))
overlay_draw.text(p_draw, "p", align="center", fill=(0, 0, 0))
if p_target is not None:
t_draw = int(p_target[0]), int(p_target[1])
overlay_draw.ellipse(
(
t_draw[0] - rad_draw,
t_draw[1] - rad_draw,
t_draw[0] + rad_draw,
t_draw[1] + rad_draw,
),
fill=t_color,
)
if curr_point is not None and curr_point == point_key:
# overlay_draw.text(t_draw, "t", font=font, align="center", fill=(0, 0, 0))
overlay_draw.text(t_draw, "t", align="center", fill=(0, 0, 0))
return Image.alpha_composite(image.convert("RGBA"),
overlay_rgba).convert("RGB")
def draw_mask_on_image(image, mask):
im_mask = np.uint8(mask * 255)
im_mask_rgba = np.concatenate(
(
np.tile(im_mask[..., None], [1, 1, 3]),
45 * np.ones(
(im_mask.shape[0], im_mask.shape[1], 1), dtype=np.uint8),
),
axis=-1,
)
im_mask_rgba = Image.fromarray(im_mask_rgba).convert("RGBA")
return Image.alpha_composite(image.convert("RGBA"),
im_mask_rgba).convert("RGB")
def on_change_single_global_state(keys,
value,
global_state,
map_transform=None):
if map_transform is not None:
value = map_transform(value)
curr_state = global_state
if isinstance(keys, str):
last_key = keys
else:
for k in keys[:-1]:
curr_state = curr_state[k]
last_key = keys[-1]
curr_state[last_key] = value
return global_state
def get_latest_points_pair(points_dict):
if not points_dict:
return None
point_idx = list(points_dict.keys())
latest_point_idx = max(point_idx)
return latest_point_idx
================================================
FILE: gui_utils/__init__.py
================================================
# 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
================================================
FILE: gui_utils/gl_utils.py
================================================
# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto. Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.
import math
import os
import functools
import contextlib
import numpy as np
import OpenGL.GL as gl
import OpenGL.GL.ARB.texture_float
import dnnlib
#----------------------------------------------------------------------------
def init_egl():
assert os.environ['PYOPENGL_PLATFORM'] == 'egl' # Must be set before importing OpenGL.
import OpenGL.EGL as egl
import ctypes
# Initialize EGL.
display = egl.eglGetDisplay(egl.EGL_DEFAULT_DISPLAY)
assert display != egl.EGL_NO_DISPLAY
major = ctypes.c_int32()
minor = ctypes.c_int32()
ok = egl.eglInitialize(display, major, minor)
assert ok
assert major.value * 10 + minor.value >= 14
# Choose config.
config_attribs = [
egl.EGL_RENDERABLE_TYPE, egl.EGL_OPENGL_BIT,
egl.EGL_SURFACE_TYPE, egl.EGL_PBUFFER_BIT,
egl.EGL_NONE
]
configs = (ctypes.c_int32 * 1)()
num_configs = ctypes.c_int32()
ok = egl.eglChooseConfig(display, config_attribs, configs, 1, num_configs)
assert ok
assert num_configs.value == 1
config = configs[0]
# Create dummy pbuffer surface.
surface_attribs = [
egl.EGL_WIDTH, 1,
egl.EGL_HEIGHT, 1,
egl.EGL_NONE
]
surface = egl.eglCreatePbufferSurface(display, config, surface_attribs)
assert surface != egl.EGL_NO_SURFACE
# Setup GL context.
ok = egl.eglBindAPI(egl.EGL_OPENGL_API)
assert ok
context = egl.eglCreateContext(display, config, egl.EGL_NO_CONTEXT, None)
assert context != egl.EGL_NO_CONTEXT
ok = egl.eglMakeCurrent(display, surface, surface, context)
assert ok
#----------------------------------------------------------------------------
_texture_formats = {
('uint8', 1): dnnlib.EasyDict(type=gl.GL_UNSIGNED_BYTE, format=gl.GL_LUMINANCE, internalformat=gl.GL_LUMINANCE8),
('uint8', 2): dnnlib.EasyDict(type=gl.GL_UNSIGNED_BYTE, format=gl.GL_LUMINANCE_ALPHA, internalformat=gl.GL_LUMINANCE8_ALPHA8),
('uint8', 3): dnnlib.EasyDict(type=gl.GL_UNSIGNED_BYTE, format=gl.GL_RGB, internalformat=gl.GL_RGB8),
('uint8', 4): dnnlib.EasyDict(type=gl.GL_UNSIGNED_BYTE, format=gl.GL_RGBA, internalformat=gl.GL_RGBA8),
('float32', 1): dnnlib.EasyDict(type=gl.GL_FLOAT, format=gl.GL_LUMINANCE, internalformat=OpenGL.GL.ARB.texture_float.GL_LUMINANCE32F_ARB),
('float32', 2): dnnlib.EasyDict(type=gl.GL_FLOAT, format=gl.GL_LUMINANCE_ALPHA, internalformat=OpenGL.GL.ARB.texture_float.GL_LUMINANCE_ALPHA32F_ARB),
('float32', 3): dnnlib.EasyDict(type=gl.GL_FLOAT, format=gl.GL_RGB, internalformat=gl.GL_RGB32F),
('float32', 4): dnnlib.EasyDict(type=gl.GL_FLOAT, format=gl.GL_RGBA, internalformat=gl.GL_RGBA32F),
}
def get_texture_format(dtype, channels):
return _texture_formats[(np.dtype(dtype).name, int(channels))]
#----------------------------------------------------------------------------
def prepare_texture_data(image):
image = np.asarray(image)
if image.ndim == 2:
image = image[:, :, np.newaxis]
if image.dtype.name == 'float64':
image = image.astype('float32')
return image
#----------------------------------------------------------------------------
def draw_pixels(image, *, pos=0, zoom=1, align=0, rint=True):
pos = np.broadcast_to(np.asarray(pos, dtype='float32'), [2])
zoom = np.broadcast_to(np.asarray(zoom, dtype='float32'), [2])
align = np.broadcast_to(np.asarray(align, dtype='float32'), [2])
image = prepare_texture_data(image)
height, width, channels = image.shape
size = zoom * [width, height]
pos = pos - size * align
if rint:
pos = np.rint(pos)
fmt = get_texture_format(image.dtype, channels)
gl.glPushAttrib(gl.GL_CURRENT_BIT | gl.GL_PIXEL_MODE_BIT)
gl.glPushClientAttrib(gl.GL_CLIENT_PIXEL_STORE_BIT)
gl.glRasterPos2f(pos[0], pos[1])
gl.glPixelZoom(zoom[0], -zoom[1])
gl.glPixelStorei(gl.GL_UNPACK_ALIGNMENT, 1)
gl.glDrawPixels(width, height, fmt.format, fmt.type, image)
gl.glPopClientAttrib()
gl.glPopAttrib()
#----------------------------------------------------------------------------
def read_pixels(width, height, *, pos=0, dtype='uint8', channels=3):
pos = np.broadcast_to(np.asarray(pos, dtype='float32'), [2])
dtype = np.dtype(dtype)
fmt = get_texture_format(dtype, channels)
image = np.empty([height, width, channels], dtype=dtype)
gl.glPushClientAttrib(gl.GL_CLIENT_PIXEL_STORE_BIT)
gl.glPixelStorei(gl.GL_PACK_ALIGNMENT, 1)
gl.glReadPixels(int(np.round(pos[0])), int(np.round(pos[1])), width, height, fmt.format, fmt.type, image)
gl.glPopClientAttrib()
return np.flipud(image)
#----------------------------------------------------------------------------
class Texture:
def __init__(self, *, image=None, width=None, height=None, channels=None, dtype=None, bilinear=True, mipmap=True):
self.gl_id = None
self.bilinear = bilinear
self.mipmap = mipmap
# Determine size and dtype.
if image is not None:
image = prepare_texture_data(image)
self.height, self.width, self.channels = image.shape
self.dtype = image.dtype
else:
assert width is not None and height is not None
self.width = width
self.height = height
self.channels = channels if channels is not None else 3
self.dtype = np.dtype(dtype) if dtype is not None else np.uint8
# Validate size and dtype.
assert isinstance(self.width, int) and self.width >= 0
assert isinstance(self.height, int) and self.height >= 0
assert isinstance(self.channels, int) and self.channels >= 1
assert self.is_compatible(width=width, height=height, channels=channels, dtype=dtype)
# Create texture object.
self.gl_id = gl.glGenTextures(1)
with self.bind():
gl.glTexParameterf(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_WRAP_S, gl.GL_CLAMP_TO_EDGE)
gl.glTexParameterf(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_WRAP_T, gl.GL_CLAMP_TO_EDGE)
gl.glTexParameterf(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_MAG_FILTER, gl.GL_LINEAR if self.bilinear else gl.GL_NEAREST)
gl.glTexParameterf(gl.GL_TEXTURE_2D, gl.GL_TEXTURE_MIN_FILTER, gl.GL_LINEAR_MIPMAP_LINEAR if self.mipmap else gl.GL_NEAREST)
self.update(image)
def delete(self):
if self.gl_id is not None:
gl.glDeleteTextures([self.gl_id])
self.gl_id = None
def __del__(self):
try:
self.delete()
except:
pass
@contextlib.contextmanager
def bind(self):
prev_id = gl.glGetInteger(gl.GL_TEXTURE_BINDING_2D)
gl.glBindTexture(gl.GL_TEXTURE_2D, self.gl_id)
yield
gl.glBindTexture(gl.GL_TEXTURE_2D, prev_id)
def update(self, image):
if image is not None:
image = prepare_texture_data(image)
assert self.is_compatible(image=image)
with self.bind():
fmt = get_texture_format(self.dtype, self.channels)
gl.glPushClientAttrib(gl.GL_CLIENT_PIXEL_STORE_BIT)
gl.glPixelStorei(gl.GL_UNPACK_ALIGNMENT, 1)
gl.glTexImage2D(gl.GL_TEXTURE_2D, 0, fmt.internalformat, self.width, self.height, 0, fmt.format, fmt.type, image)
if self.mipmap:
gl.glGenerateMipmap(gl.GL_TEXTURE_2D)
gl.glPopClientAttrib()
def draw(self, *, pos=0, zoom=1, align=0, rint=False, color=1, alpha=1, rounding=0):
zoom = np.broadcast_to(np.asarray(zoom, dtype='float32'), [2])
size = zoom * [self.width, self.height]
with self.bind():
gl.glPushAttrib(gl.GL_ENABLE_BIT)
gl.glEnable(gl.GL_TEXTURE_2D)
draw_rect(pos=pos, size=size, align=align, rint=rint, color=color, alpha=alpha, rounding=rounding)
gl.glPopAttrib()
def is_compatible(self, *, image=None, width=None, height=None, channels=None, dtype=None): # pylint: disable=too-many-return-statements
if image is not None:
if image.ndim != 3:
return False
ih, iw, ic = image.shape
if not self.is_compatible(width=iw, height=ih, channels=ic, dtype=image.dtype):
return False
if width is not None and self.width != width:
return False
if height is not None and self.height != height:
return False
if channels is not None and self.channels != channels:
return False
if dtype is not None and self.dtype != dtype:
return False
return True
#----------------------------------------------------------------------------
class Framebuffer:
def __init__(self, *, texture=None, width=None, height=None, channels=None, dtype=None, msaa=0):
self.texture = texture
self.gl_id = None
self.gl_color = None
self.gl_depth_stencil = None
self.msaa = msaa
# Determine size and dtype.
if texture is not None:
assert isinstance(self.texture, Texture)
self.width = texture.width
self.height = texture.height
self.channels = texture.channels
self.dtype = texture.dtype
else:
assert width is not None and height is not None
self.width = width
self.height = height
self.channels = channels if channels is not None else 4
self.dtype = np.dtype(dtype) if dtype is not None else np.float32
# Validate size and dtype.
assert isinstance(self.width, int) and self.width >= 0
assert isinstance(self.height, int) and self.height >= 0
assert isinstance(self.channels, int) and self.channels >= 1
assert width is None or width == self.width
assert height is None or height == self.height
assert channels is None or channels == self.channels
assert dtype is None or dtype == self.dtype
# Create framebuffer object.
self.gl_id = gl.glGenFramebuffers(1)
with self.bind():
# Setup color buffer.
if self.texture is not None:
assert self.msaa == 0
gl.glFramebufferTexture2D(gl.GL_FRAMEBUFFER, gl.GL_COLOR_ATTACHMENT0, gl.GL_TEXTURE_2D, self.texture.gl_id, 0)
else:
fmt = get_texture_format(self.dtype, self.channels)
self.gl_color = gl.glGenRenderbuffers(1)
gl.glBindRenderbuffer(gl.GL_RENDERBUFFER, self.gl_color)
gl.glRenderbufferStorageMultisample(gl.GL_RENDERBUFFER, self.msaa, fmt.internalformat, self.width, self.height)
gl.glFramebufferRenderbuffer(gl.GL_FRAMEBUFFER, gl.GL_COLOR_ATTACHMENT0, gl.GL_RENDERBUFFER, self.gl_color)
# Setup depth/stencil buffer.
self.gl_depth_stencil = gl.glGenRenderbuffers(1)
gl.glBindRenderbuffer(gl.GL_RENDERBUFFER, self.gl_depth_stencil)
gl.glRenderbufferStorageMultisample(gl.GL_RENDERBUFFER, self.msaa, gl.GL_DEPTH24_STENCIL8, self.width, self.height)
gl.glFramebufferRenderbuffer(gl.GL_FRAMEBUFFER, gl.GL_DEPTH_STENCIL_ATTACHMENT, gl.GL_RENDERBUFFER, self.gl_depth_stencil)
def delete(self):
if self.gl_id is not None:
gl.glDeleteFramebuffers([self.gl_id])
self.gl_id = None
if self.gl_color is not None:
gl.glDeleteRenderbuffers(1, [self.gl_color])
self.gl_color = None
if self.gl_depth_stencil is not None:
gl.glDeleteRenderbuffers(1, [self.gl_depth_stencil])
self.gl_depth_stencil = None
def __del__(self):
try:
self.delete()
except:
pass
@contextlib.contextmanager
def bind(self):
prev_fbo = gl.glGetInteger(gl.GL_FRAMEBUFFER_BINDING)
prev_rbo = gl.glGetInteger(gl.GL_RENDERBUFFER_BINDING)
gl.glBindFramebuffer(gl.GL_FRAMEBUFFER, self.gl_id)
if self.width is not None and self.height is not None:
gl.glViewport(0, 0, self.width, self.height)
yield
gl.glBindFramebuffer(gl.GL_FRAMEBUFFER, prev_fbo)
gl.glBindRenderbuffer(gl.GL_RENDERBUFFER, prev_rbo)
def blit(self, dst=None):
assert dst is None or isinstance(dst, Framebuffer)
with self.bind():
gl.glBindFramebuffer(gl.GL_DRAW_FRAMEBUFFER, 0 if dst is None else dst.fbo)
gl.glBlitFramebuffer(0, 0, self.width, self.height, 0, 0, self.width, self.height, gl.GL_COLOR_BUFFER_BIT, gl.GL_NEAREST)
#----------------------------------------------------------------------------
def draw_shape(vertices, *, mode=gl.GL_TRIANGLE_FAN, pos=0, size=1, color=1, alpha=1):
assert vertices.ndim == 2 and vertices.shape[1] == 2
pos = np.broadcast_to(np.asarray(pos, dtype='float32'), [2])
size = np.broadcast_to(np.asarray(size, dtype='float32'), [2])
color = np.broadcast_to(np.asarray(color, dtype='float32'), [3])
alpha = np.clip(np.broadcast_to(np.asarray(alpha, dtype='float32'), []), 0, 1)
gl.glPushClientAttrib(gl.GL_CLIENT_VERTEX_ARRAY_BIT)
gl.glPushAttrib(gl.GL_CURRENT_BIT | gl.GL_TRANSFORM_BIT)
gl.glMatrixMode(gl.GL_MODELVIEW)
gl.glPushMatrix()
gl.glEnableClientState(gl.GL_VERTEX_ARRAY)
gl.glEnableClientState(gl.GL_TEXTURE_COORD_ARRAY)
gl.glVertexPointer(2, gl.GL_FLOAT, 0, vertices)
gl.glTexCoordPointer(2, gl.GL_FLOAT, 0, vertices)
gl.glTranslate(pos[0], pos[1], 0)
gl.glScale(size[0], size[1], 1)
gl.glColor4f(color[0] * alpha, color[1] * alpha, color[2] * alpha, alpha)
gl.glDrawArrays(mode, 0, vertices.shape[0])
gl.glPopMatrix()
gl.glPopAttrib()
gl.glPopClientAttrib()
#----------------------------------------------------------------------------
def draw_arrow(x1, y1, x2, y2, l=10, width=1.0):
# Compute the length and angle of the arrow
dx = x2 - x1
dy = y2 - y1
length = math.sqrt(dx**2 + dy**2)
if length < l:
return
angle = math.atan2(dy, dx)
# Save the current modelview matrix
gl.glPushMatrix()
# Translate and rotate the coordinate system
gl.glTranslatef(x1, y1, 0.0)
gl.glRotatef(angle * 180.0 / math.pi, 0.0, 0.0, 1.0)
# Set the line width
gl.glLineWidth(width)
# gl.glColor3f(0.75, 0.75, 0.75)
# Begin drawing lines
gl.glBegin(gl.GL_LINES)
# Draw the shaft of the arrow
gl.glVertex2f(0.0, 0.0)
gl.glVertex2f(length, 0.0)
# Draw the head of the arrow
gl.glVertex2f(length, 0.0)
gl.glVertex2f(length - 2 * l, l)
gl.glVertex2f(length, 0.0)
gl.glVertex2f(length - 2 * l, -l)
# End drawing lines
gl.glEnd()
# Restore the modelview matrix
gl.glPopMatrix()
#----------------------------------------------------------------------------
def draw_rect(*, pos=0, pos2=None, size=None, align=0, rint=False, color=1, alpha=1, rounding=0):
assert pos2 is None or size is None
pos = np.broadcast_to(np.asarray(pos, dtype='float32'), [2])
pos2 = np.broadcast_to(np.asarray(pos2, dtype='float32'), [2]) if pos2 is not None else None
size = np.broadcast_to(np.asarray(size, dtype='float32'), [2]) if size is not None else None
size = size if size is not None else pos2 - pos if pos2 is not None else np.array([1, 1], dtype='float32')
pos = pos - size * align
if rint:
pos = np.rint(pos)
rounding = np.broadcast_to(np.asarray(rounding, dtype='float32'), [2])
rounding = np.minimum(np.abs(rounding) / np.maximum(np.abs(size), 1e-8), 0.5)
if np.min(rounding) == 0:
rounding *= 0
vertices = _setup_rect(float(rounding[0]), float(rounding[1]))
draw_shape(vertices, mode=gl.GL_TRIANGLE_FAN, pos=pos, size=size, color=color, alpha=alpha)
@functools.lru_cache(maxsize=10000)
def _setup_rect(rx, ry):
t = np.linspace(0, np.pi / 2, 1 if max(rx, ry) == 0 else 64)
s = 1 - np.sin(t); c = 1 - np.cos(t)
x = [c * rx, 1 - s * rx, 1 - c * rx, s * rx]
y = [s * ry, c * ry, 1 - s * ry, 1 - c * ry]
v = np.stack([x, y], axis=-1).reshape(-1, 2)
return v.astype('float32')
#----------------------------------------------------------------------------
def draw_circle(*, center=0, radius=100, hole=0, color=1, alpha=1):
hole = np.broadcast_to(np.asarray(hole, dtype='float32'), [])
vertices = _setup_circle(float(hole))
draw_shape(vertices, mode=gl.GL_TRIANGLE_STRIP, pos=center, size=radius, color=color, alpha=alpha)
@functools.lru_cache(maxsize=10000)
def _setup_circle(hole):
t = np.linspace(0, np.pi * 2, 128)
s = np.sin(t); c = np.cos(t)
v = np.stack([c, s, c * hole, s * hole], axis=-1).reshape(-1, 2)
return v.astype('float32')
#----------------------------------------------------------------------------
================================================
FILE: gui_utils/glfw_window.py
================================================
# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto. Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.
import time
import glfw
import OpenGL.GL as gl
from . import gl_utils
#----------------------------------------------------------------------------
class GlfwWindow: # pylint: disable=too-many-public-methods
def __init__(self, *, title='GlfwWindow', window_width=1920, window_height=1080, deferred_show=True, close_on_esc=True):
self._glfw_window = None
self._drawing_frame = False
self._frame_start_time = None
self._frame_delta = 0
self._fps_limit = None
self._vsync = None
self._skip_frames = 0
self._deferred_show = deferred_show
self._close_on_esc = close_on_esc
self._esc_pressed = False
self._drag_and_drop_paths = None
self._capture_next_frame = False
self._captured_frame = None
# Create window.
glfw.init()
glfw.window_hint(glfw.VISIBLE, False)
self._glfw_window = glfw.create_window(width=window_width, height=window_height, title=title, monitor=None, share=None)
self._attach_glfw_callbacks()
self.make_context_current()
# Adjust window.
self.set_vsync(False)
self.set_window_size(window_width, window_height)
if not self._deferred_show:
glfw.show_window(self._glfw_window)
def close(self):
if self._drawing_frame:
self.end_frame()
if self._glfw_window is not None:
glfw.destroy_window(self._glfw_window)
self._glfw_window = None
#glfw.terminate() # Commented out to play it nice with other glfw clients.
def __del__(self):
try:
self.close()
except:
pass
@property
def window_width(self):
return self.content_width
@property
def window_height(self):
return self.content_height + self.title_bar_height
@property
def content_width(self):
width, _height = glfw.get_window_size(self._glfw_window)
return width
@property
def content_height(self):
_width, height = glfw.get_window_size(self._glfw_window)
return height
@property
def title_bar_height(self):
_left, top, _right, _bottom = glfw.get_window_frame_size(self._glfw_window)
return top
@property
def monitor_width(self):
_, _, width, _height = glfw.get_monitor_workarea(glfw.get_primary_monitor())
return width
@property
def monitor_height(self):
_, _, _width, height = glfw.get_monitor_workarea(glfw.get_primary_monitor())
return height
@property
def frame_delta(self):
return self._frame_delta
def set_title(self, title):
glfw.set_window_title(self._glfw_window, title)
def set_window_size(self, width, height):
width = min(width, self.monitor_width)
height = min(height, self.monitor_height)
glfw.set_window_size(self._glfw_window, width, max(height - self.title_bar_height, 0))
if width == self.monitor_width and height == self.monitor_height:
self.maximize()
def set_content_size(self, width, height):
self.set_window_size(width, height + self.title_bar_height)
def maximize(self):
glfw.maximize_window(self._glfw_window)
def set_position(self, x, y):
glfw.set_window_pos(self._glfw_window, x, y + self.title_bar_height)
def center(self):
self.set_position((self.monitor_width - self.window_width) // 2, (self.monitor_height - self.window_height) // 2)
def set_vsync(self, vsync):
vsync = bool(vsync)
if vsync != self._vsync:
glfw.swap_interval(1 if vsync else 0)
self._vsync = vsync
def set_fps_limit(self, fps_limit):
self._fps_limit = int(fps_limit)
def should_close(self):
return glfw.window_should_close(self._glfw_window) or (self._close_on_esc and self._esc_pressed)
def skip_frame(self):
self.skip_frames(1)
def skip_frames(self, num): # Do not update window for the next N frames.
self._skip_frames = max(self._skip_frames, int(num))
def is_skipping_frames(self):
return self._skip_frames > 0
def capture_next_frame(self):
self._capture_next_frame = True
def pop_captured_frame(self):
frame = self._captured_frame
self._captured_frame = None
return frame
def pop_drag_and_drop_paths(self):
paths = self._drag_and_drop_paths
self._drag_and_drop_paths = None
return paths
def draw_frame(self): # To be overridden by subclass.
self.begin_frame()
# Rendering code goes here.
self.end_frame()
def make_context_current(self):
if self._glfw_window is not None:
glfw.make_context_current(self._glfw_window)
def begin_frame(self):
# End previous frame.
if self._drawing_frame:
self.end_frame()
# Apply FPS limit.
if self._frame_start_time is not None and self._fps_limit is not None:
delay = self._frame_start_time - time.perf_counter() + 1 / self._fps_limit
if delay > 0:
time.sleep(delay)
cur_time = time.perf_counter()
if self._frame_start_time is not None:
self._frame_delta = cur_time - self._frame_start_time
self._frame_start_time = cur_time
# Process events.
glfw.poll_events()
# Begin frame.
self._drawing_frame = True
self.make_context_current()
# Initialize GL state.
gl.glViewport(0, 0, self.content_width, self.content_height)
gl.glMatrixMode(gl.GL_PROJECTION)
gl.glLoadIdentity()
gl.glTranslate(-1, 1, 0)
gl.glScale(2 / max(self.content_width, 1), -2 / max(self.content_height, 1), 1)
gl.glMatrixMode(gl.GL_MODELVIEW)
gl.glLoadIdentity()
gl.glEnable(gl.GL_BLEND)
gl.glBlendFunc(gl.GL_ONE, gl.GL_ONE_MINUS_SRC_ALPHA) # Pre-multiplied alpha.
# Clear.
gl.glClearColor(0, 0, 0, 1)
gl.glClear(gl.GL_COLOR_BUFFER_BIT | gl.GL_DEPTH_BUFFER_BIT)
def end_frame(self):
assert self._drawing_frame
self._drawing_frame = False
# Skip frames if requested.
if self._skip_frames > 0:
self._skip_frames -= 1
return
# Capture frame if requested.
if self._capture_next_frame:
self._captured_frame = gl_utils.read_pixels(self.content_width, self.content_height)
self._capture_next_frame = False
# Update window.
if self._deferred_show:
glfw.show_window(self._glfw_window)
self._deferred_show = False
glfw.swap_buffers(self._glfw_window)
def _attach_glfw_callbacks(self):
glfw.set_key_callback(self._glfw_window, self._glfw_key_callback)
glfw.set_drop_callback(self._glfw_window, self._glfw_drop_callback)
def _glfw_key_callback(self, _window, key, _scancode, action, _mods):
if action == glfw.PRESS and key == glfw.KEY_ESCAPE:
self._esc_pressed = True
def _glfw_drop_callback(self, _window, paths):
self._drag_and_drop_paths = paths
#----------------------------------------------------------------------------
================================================
FILE: gui_utils/imgui_utils.py
================================================
# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto. Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.
import contextlib
import imgui
#----------------------------------------------------------------------------
def set_default_style(color_scheme='dark', spacing=9, indent=23, scrollbar=27):
s = imgui.get_style()
s.window_padding = [spacing, spacing]
s.item_spacing = [spacing, spacing]
s.item_inner_spacing = [spacing, spacing]
s.columns_min_spacing = spacing
s.indent_spacing = indent
s.scrollbar_size = scrollbar
s.frame_padding = [4, 3]
s.window_border_size = 1
s.child_border_size = 1
s.popup_border_size = 1
s.frame_border_size = 1
s.window_rounding = 0
s.child_rounding = 0
s.popup_rounding = 3
s.frame_rounding = 3
s.scrollbar_rounding = 3
s.grab_rounding = 3
getattr(imgui, f'style_colors_{color_scheme}')(s)
c0 = s.colors[imgui.COLOR_MENUBAR_BACKGROUND]
c1 = s.colors[imgui.COLOR_FRAME_BACKGROUND]
s.colors[imgui.COLOR_POPUP_BACKGROUND] = [x * 0.7 + y * 0.3 for x, y in zip(c0, c1)][:3] + [1]
#----------------------------------------------------------------------------
@contextlib.contextmanager
def grayed_out(cond=True):
if cond:
s = imgui.get_style()
text = s.colors[imgui.COLOR_TEXT_DISABLED]
grab = s.colors[imgui.COLOR_SCROLLBAR_GRAB]
back = s.colors[imgui.COLOR_MENUBAR_BACKGROUND]
imgui.push_style_color(imgui.COLOR_TEXT, *text)
imgui.push_style_color(imgui.COLOR_CHECK_MARK, *grab)
imgui.push_style_color(imgui.COLOR_SLIDER_GRAB, *grab)
imgui.push_style_color(imgui.COLOR_SLIDER_GRAB_ACTIVE, *grab)
imgui.push_style_color(imgui.COLOR_FRAME_BACKGROUND, *back)
imgui.push_style_color(imgui.COLOR_FRAME_BACKGROUND_HOVERED, *back)
imgui.push_style_color(imgui.COLOR_FRAME_BACKGROUND_ACTIVE, *back)
imgui.push_style_color(imgui.COLOR_BUTTON, *back)
imgui.push_style_color(imgui.COLOR_BUTTON_HOVERED, *back)
imgui.push_style_color(imgui.COLOR_BUTTON_ACTIVE, *back)
imgui.push_style_color(imgui.COLOR_HEADER, *back)
imgui.push_style_color(imgui.COLOR_HEADER_HOVERED, *back)
imgui.push_style_color(imgui.COLOR_HEADER_ACTIVE, *back)
imgui.push_style_color(imgui.COLOR_POPUP_BACKGROUND, *back)
yield
imgui.pop_style_color(14)
else:
yield
#----------------------------------------------------------------------------
@contextlib.contextmanager
def item_width(width=None):
if width is not None:
imgui.push_item_width(width)
yield
imgui.pop_item_width()
else:
yield
#----------------------------------------------------------------------------
def scoped_by_object_id(method):
def decorator(self, *args, **kwargs):
imgui.push_id(str(id(self)))
res = method(self, *args, **kwargs)
imgui.pop_id()
return res
return decorator
#----------------------------------------------------------------------------
def button(label, width=0, enabled=True):
with grayed_out(not enabled):
clicked = imgui.button(label, width=width)
clicked = clicked and enabled
return clicked
#----------------------------------------------------------------------------
def collapsing_header(text, visible=None, flags=0, default=False, enabled=True, show=True):
expanded = False
if show:
if default:
flags |= imgui.TREE_NODE_DEFAULT_OPEN
if not enabled:
flags |= imgui.TREE_NODE_LEAF
with grayed_out(not enabled):
expanded, visible = imgui.collapsing_header(text, visible=visible, flags=flags)
expanded = expanded and enabled
return expanded, visible
#----------------------------------------------------------------------------
def popup_button(label, width=0, enabled=True):
if button(label, width, enabled):
imgui.open_popup(label)
opened = imgui.begin_popup(label)
return opened
#----------------------------------------------------------------------------
def input_text(label, value, buffer_length, flags, width=None, help_text=''):
old_value = value
color = list(imgui.get_style().colors[imgui.COLOR_TEXT])
if value == '':
color[-1] *= 0.5
with item_width(width):
imgui.push_style_color(imgui.COLOR_TEXT, *color)
value = value if value != '' else help_text
changed, value = imgui.input_text(label, value, buffer_length, flags)
value = value if value != help_text else ''
imgui.pop_style_color(1)
if not flags & imgui.INPUT_TEXT_ENTER_RETURNS_TRUE:
changed = (value != old_value)
return changed, value
#----------------------------------------------------------------------------
def drag_previous_control(enabled=True):
dragging = False
dx = 0
dy = 0
if imgui.begin_drag_drop_source(imgui.DRAG_DROP_SOURCE_NO_PREVIEW_TOOLTIP):
if enabled:
dragging = True
dx, dy = imgui.get_mouse_drag_delta()
imgui.reset_mouse_drag_delta()
imgui.end_drag_drop_source()
return dragging, dx, dy
#----------------------------------------------------------------------------
def drag_button(label, width=0, enabled=True):
clicked = button(label, width=width, enabled=enabled)
dragging, dx, dy = drag_previous_control(enabled=enabled)
return clicked, dragging, dx, dy
#----------------------------------------------------------------------------
def drag_hidden_window(label, x, y, width, height, enabled=True):
imgui.push_style_color(imgui.COLOR_WINDOW_BACKGROUND, 0, 0, 0, 0)
imgui.push_style_color(imgui.COLOR_BORDER, 0, 0, 0, 0)
imgui.set_next_window_position(x, y)
imgui.set_next_window_size(width, height)
imgui.begin(label, closable=False, flags=(imgui.WINDOW_NO_TITLE_BAR | imgui.WINDOW_NO_RESIZE | imgui.WINDOW_NO_MOVE))
dragging, dx, dy = drag_previous_control(enabled=enabled)
imgui.end()
imgui.pop_style_color(2)
return dragging, dx, dy
#----------------------------------------------------------------------------
def click_hidden_window(label, x, y, width, height, img_w, img_h, enabled=True):
imgui.push_style_color(imgui.COLOR_WINDOW_BACKGROUND, 0, 0, 0, 0)
imgui.push_style_color(imgui.COLOR_BORDER, 0, 0, 0, 0)
imgui.set_next_window_position(x, y)
imgui.set_next_window_size(width, height)
imgui.begin(label, closable=False, flags=(imgui.WINDOW_NO_TITLE_BAR | imgui.WINDOW_NO_RESIZE | imgui.WINDOW_NO_MOVE))
clicked, down = False, False
img_x, img_y = 0, 0
if imgui.is_mouse_down():
posx, posy = imgui.get_mouse_pos()
if posx >= x and posx < x + width and posy >= y and posy < y + height:
if imgui.is_mouse_clicked():
clicked = True
down = True
img_x = round((posx - x) / (width - 1) * (img_w - 1))
img_y = round((posy - y) / (height - 1) * (img_h - 1))
imgui.end()
imgui.pop_style_color(2)
return clicked, down, img_x, img_y
#----------------------------------------------------------------------------
================================================
FILE: gui_utils/imgui_window.py
================================================
# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto. Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.
import os
import imgui
import imgui.integrations.glfw
from . import glfw_window
from . import imgui_utils
from . import text_utils
#----------------------------------------------------------------------------
class ImguiWindow(glfw_window.GlfwWindow):
def __init__(self, *, title='ImguiWindow', font=None, font_sizes=range(14,24), **glfw_kwargs):
if font is None:
font = text_utils.get_default_font()
font_sizes = {int(size) for size in font_sizes}
super().__init__(title=title, **glfw_kwargs)
# Init fields.
self._imgui_context = None
self._imgui_renderer = None
self._imgui_fonts = None
self._cur_font_size = max(font_sizes)
# Delete leftover imgui.ini to avoid unexpected behavior.
if os.path.isfile('imgui.ini'):
os.remove('imgui.ini')
# Init ImGui.
self._imgui_context = imgui.create_context()
self._imgui_renderer = _GlfwRenderer(self._glfw_window)
self._attach_glfw_callbacks()
imgui.get_io().ini_saving_rate = 0 # Disable creating imgui.ini at runtime.
imgui.get_io().mouse_drag_threshold = 0 # Improve behavior with imgui_utils.drag_custom().
self._imgui_fonts = {size: imgui.get_io().fonts.add_font_from_file_ttf(font, size) for size in font_sizes}
self._imgui_renderer.refresh_font_texture()
def close(self):
self.make_context_current()
self._imgui_fonts = None
if self._imgui_renderer is not None:
self._imgui_renderer.shutdown()
self._imgui_renderer = None
if self._imgui_context is not None:
#imgui.destroy_context(self._imgui_context) # Commented out to avoid creating imgui.ini at the end.
self._imgui_context = None
super().close()
def _glfw_key_callback(self, *args):
super()._glfw_key_callback(*args)
self._imgui_renderer.keyboard_callback(*args)
@property
def font_size(self):
return self._cur_font_size
@property
def spacing(self):
return round(self._cur_font_size * 0.4)
def set_font_size(self, target): # Applied on next frame.
self._cur_font_size = min((abs(key - target), key) for key in self._imgui_fonts.keys())[1]
def begin_frame(self):
# Begin glfw frame.
super().begin_frame()
# Process imgui events.
self._imgui_renderer.mouse_wheel_multiplier = self._cur_font_size / 10
if self.content_width > 0 and self.content_height > 0:
self._imgui_renderer.process_inputs()
# Begin imgui frame.
imgui.new_frame()
imgui.push_font(self._imgui_fonts[self._cur_font_size])
imgui_utils.set_default_style(spacing=self.spacing, indent=self.font_size, scrollbar=self.font_size+4)
def end_frame(self):
imgui.pop_font()
imgui.render()
imgui.end_frame()
self._imgui_renderer.render(imgui.get_draw_data())
super().end_frame()
#----------------------------------------------------------------------------
# Wrapper class for GlfwRenderer to fix a mouse wheel bug on Linux.
class _GlfwRenderer(imgui.integrations.glfw.GlfwRenderer):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.mouse_wheel_multiplier = 1
def scroll_callback(self, window, x_offset, y_offset):
self.io.mouse_wheel += y_offset * self.mouse_wheel_multiplier
#----------------------------------------------------------------------------
================================================
FILE: gui_utils/text_utils.py
================================================
# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto. Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.
import functools
from typing import Optional
import dnnlib
import numpy as np
import PIL.Image
import PIL.ImageFont
import scipy.ndimage
from . import gl_utils
#----------------------------------------------------------------------------
def get_default_font():
url = 'http://fonts.gstatic.com/s/opensans/v17/mem8YaGs126MiZpBA-U1UpcaXcl0Aw.ttf' # Open Sans regular
return dnnlib.util.open_url(url, return_filename=True)
#----------------------------------------------------------------------------
@functools.lru_cache(maxsize=None)
def get_pil_font(font=None, size=32):
if font is None:
font = get_default_font()
return PIL.ImageFont.truetype(font=font, size=size)
#----------------------------------------------------------------------------
def get_array(string, *, dropshadow_radius: int=None, **kwargs):
if dropshadow_radius is not None:
offset_x = int(np.ceil(dropshadow_radius*2/3))
offset_y = int(np.ceil(dropshadow_radius*2/3))
return _get_array_priv(string, dropshadow_radius=dropshadow_radius, offset_x=offset_x, offset_y=offset_y, **kwargs)
else:
return _get_array_priv(string, **kwargs)
@functools.lru_cache(maxsize=10000)
def _get_array_priv(
string: str, *,
size: int = 32,
max_width: Optional[int]=None,
max_height: Optional[int]=None,
min_size=10,
shrink_coef=0.8,
dropshadow_radius: int=None,
offset_x: int=None,
offset_y: int=None,
**kwargs
):
cur_size = size
array = None
while True:
if dropshadow_radius is not None:
# separate implementation for dropshadow text rendering
array = _get_array_impl_dropshadow(string, size=cur_size, radius=dropshadow_radius, offset_x=offset_x, offset_y=offset_y, **kwargs)
else:
array = _get_array_impl(string, size=cur_size, **kwargs)
height, width, _ = array.shape
if (max_width is None or width <= max_width) and (max_height is None or height <= max_height) or (cur_size <= min_size):
break
cur_size = max(int(cur_size * shrink_coef), min_size)
return array
#----------------------------------------------------------------------------
@functools.lru_cache(maxsize=10000)
def _get_array_impl(string, *, font=None, size=32, outline=0, outline_pad=3, outline_coef=3, outline_exp=2, line_pad: int=None):
pil_font = get_pil_font(font=font, size=size)
lines = [pil_font.getmask(line, 'L') for line in string.split('\n')]
lines = [np.array(line, dtype=np.uint8).reshape([line.size[1], line.size[0]]) for line in lines]
width = max(line.shape[1] for line in lines)
lines = [np.pad(line, ((0, 0), (0, width - line.shape[1])), mode='constant') for line in lines]
line_spacing = line_pad if line_pad is not None else size // 2
lines = [np.pad(line, ((0, line_spacing), (0, 0)), mode='constant') for line in lines[:-1]] + lines[-1:]
mask = np.concatenate(lines, axis=0)
alpha = mask
if outline > 0:
mask = np.pad(mask, int(np.ceil(outline * outline_pad)), mode='constant', constant_values=0)
alpha = mask.astype(np.float32) / 255
alpha = scipy.ndimage.gaussian_filter(alpha, outline)
alpha = 1 - np.maximum(1 - alpha * outline_coef, 0) ** outline_exp
alpha = (alpha * 255 + 0.5).clip(0, 255).astype(np.uint8)
alpha = np.maximum(alpha, mask)
return np.stack([mask, alpha], axis=-1)
#----------------------------------------------------------------------------
@functools.lru_cache(maxsize=10000)
def _get_array_impl_dropshadow(string, *, font=None, size=32, radius: int, offset_x: int, offset_y: int, line_pad: int=None, **kwargs):
assert (offset_x > 0) and (offset_y > 0)
pil_font = get_pil_font(font=font, size=size)
lines = [pil_font.getmask(line, 'L') for line in string.split('\n')]
lines = [np.array(line, dtype=np.uint8).reshape([line.size[1], line.size[0]]) for line in lines]
width = max(line.shape[1] for line in lines)
lines = [np.pad(line, ((0, 0), (0, width - line.shape[1])), mode='constant') for line in lines]
line_spacing = line_pad if line_pad is not None else size // 2
lines = [np.pad(line, ((0, line_spacing), (0, 0)), mode='constant') for line in lines[:-1]] + lines[-1:]
mask = np.concatenate(lines, axis=0)
alpha = mask
mask = np.pad(mask, 2*radius + max(abs(offset_x), abs(offset_y)), mode='constant', constant_values=0)
alpha = mask.astype(np.float32) / 255
alpha = scipy.ndimage.gaussian_filter(alpha, radius)
alpha = 1 - np.maximum(1 - alpha * 1.5, 0) ** 1.4
alpha = (alpha * 255 + 0.5).clip(0, 255).astype(np.uint8)
alpha = np.pad(alpha, [(offset_y, 0), (offset_x, 0)], mode='constant')[:-offset_y, :-offset_x]
alpha = np.maximum(alpha, mask)
return np.stack([mask, alpha], axis=-1)
#----------------------------------------------------------------------------
@functools.lru_cache(maxsize=10000)
def get_texture(string, bilinear=True, mipmap=True, **kwargs):
return gl_utils.Texture(image=get_array(string, **kwargs), bilinear=bilinear, mipmap=mipmap)
#----------------------------------------------------------------------------
================================================
FILE: legacy.py
================================================
# 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.
"""Converting legacy network pickle into the new format."""
import click
import pickle
import re
import copy
import numpy as np
import torch
import dnnlib
from torch_utils import misc
#----------------------------------------------------------------------------
def load_network_pkl(f, force_fp16=False):
data = _LegacyUnpickler(f).load()
# Legacy TensorFlow pickle => convert.
if isinstance(data, tuple) and len(data) == 3 and all(isinstance(net, _TFNetworkStub) for net in data):
tf_G, tf_D, tf_Gs = data
G = convert_tf_generator(tf_G)
D = convert_tf_discriminator(tf_D)
G_ema = convert_tf_generator(tf_Gs)
data = dict(G=G, D=D, G_ema=G_ema)
# Add missing fields.
if 'training_set_kwargs' not in data:
data['training_set_kwargs'] = None
if 'augment_pipe' not in data:
data['augment_pipe'] = None
# Validate contents.
assert isinstance(data['G'], torch.nn.Module)
assert isinstance(data['D'], torch.nn.Module)
assert isinstance(data['G_ema'], torch.nn.Module)
assert isinstance(data['training_set_kwargs'], (dict, type(None)))
assert isinstance(data['augment_pipe'], (torch.nn.Module, type(None)))
# Force FP16.
if force_fp16:
for key in ['G', 'D', 'G_ema']:
old = data[key]
kwargs = copy.deepcopy(old.init_kwargs)
fp16_kwargs = kwargs.get('synthesis_kwargs', kwargs)
fp16_kwargs.num_fp16_res = 4
fp16_kwargs.conv_clamp = 256
if kwargs != old.init_kwargs:
new = type(old)(**kwargs).eval().requires_grad_(False)
misc.copy_params_and_buffers(old, new, require_all=True)
data[key] = new
return data
#----------------------------------------------------------------------------
class _TFNetworkStub(dnnlib.EasyDict):
pass
class _LegacyUnpickler(pickle.Unpickler):
def find_class(self, module, name):
if module == 'dnnlib.tflib.network' and name == 'Network':
return _TFNetworkStub
return super().find_class(module, name)
#----------------------------------------------------------------------------
def _collect_tf_params(tf_net):
# pylint: disable=protected-access
tf_params = dict()
def recurse(prefix, tf_net):
for name, value in tf_net.variables:
tf_params[prefix + name] = value
for name, comp in tf_net.components.items():
recurse(prefix + name + '/', comp)
recurse('', tf_net)
return tf_params
#----------------------------------------------------------------------------
def _populate_module_params(module, *patterns):
for name, tensor in misc.named_params_and_buffers(module):
found = False
value = None
for pattern, value_fn in zip(patterns[0::2], patterns[1::2]):
match = re.fullmatch(pattern, name)
if match:
found = True
if value_fn is not None:
value = value_fn(*match.groups())
break
try:
assert found
if value is not None:
tensor.copy_(torch.from_numpy(np.array(value)))
except:
print(name, list(tensor.shape))
raise
#----------------------------------------------------------------------------
def convert_tf_generator(tf_G):
if tf_G.version < 4:
raise ValueError('TensorFlow pickle version too low')
# Collect kwargs.
tf_kwargs = tf_G.static_kwargs
known_kwargs = set()
def kwarg(tf_name, default=None, none=None):
known_kwargs.add(tf_name)
val = tf_kwargs.get(tf_name, default)
return val if val is not None else none
# Convert kwargs.
from training import networks_stylegan2
network_class = networks_stylegan2.Generator
kwargs = dnnlib.EasyDict(
z_dim = kwarg('latent_size', 512),
c_dim = kwarg('label_size', 0),
w_dim = kwarg('dlatent_size', 512),
img_resolution = kwarg('resolution', 1024),
img_channels = kwarg('num_channels', 3),
channel_base = kwarg('fmap_base', 16384) * 2,
channel_max = kwarg('fmap_max', 512),
num_fp16_res = kwarg('num_fp16_res', 0),
conv_clamp = kwarg('conv_clamp', None),
architecture = kwarg('architecture', 'skip'),
resample_filter = kwarg('resample_kernel', [1,3,3,1]),
use_noise = kwarg('use_noise', True),
activation = kwarg('nonlinearity', 'lrelu'),
mapping_kwargs = dnnlib.EasyDict(
num_layers = kwarg('mapping_layers', 8),
embed_features = kwarg('label_fmaps', None),
layer_features = kwarg('mapping_fmaps', None),
activation = kwarg('mapping_nonlinearity', 'lrelu'),
lr_multiplier = kwarg('mapping_lrmul', 0.01),
w_avg_beta = kwarg('w_avg_beta', 0.995, none=1),
),
)
# Check for unknown kwargs.
kwarg('truncation_psi')
kwarg('truncation_cutoff')
kwarg('style_mixing_prob')
kwarg('structure')
kwarg('conditioning')
kwarg('fused_modconv')
unknown_kwargs = list(set(tf_kwargs.keys()) - known_kwargs)
if len(unknown_kwargs) > 0:
raise ValueError('Unknown TensorFlow kwarg', unknown_kwargs[0])
# Collect params.
tf_params = _collect_tf_params(tf_G)
for name, value in list(tf_params.items()):
match = re.fullmatch(r'ToRGB_lod(\d+)/(.*)', name)
if match:
r = kwargs.img_resolution // (2 ** int(match.group(1)))
tf_params[f'{r}x{r}/ToRGB/{match.group(2)}'] = value
kwargs.synthesis.kwargs.architecture = 'orig'
#for name, value in tf_params.items(): print(f'{name:<50s}{list(value.shape)}')
# Convert params.
G = network_class(**kwargs).eval().requires_grad_(False)
# pylint: disable=unnecessary-lambda
# pylint: disable=f-string-without-interpolation
_populate_module_params(G,
r'mapping\.w_avg', lambda: tf_params[f'dlatent_avg'],
r'mapping\.embed\.weight', lambda: tf_params[f'mapping/LabelEmbed/weight'].transpose(),
r'mapping\.embed\.bias', lambda: tf_params[f'mapping/LabelEmbed/bias'],
r'mapping\.fc(\d+)\.weight', lambda i: tf_params[f'mapping/Dense{i}/weight'].transpose(),
r'mapping\.fc(\d+)\.bias', lambda i: tf_params[f'mapping/Dense{i}/bias'],
r'synthesis\.b4\.const', lambda: tf_params[f'synthesis/4x4/Const/const'][0],
r'synthesis\.b4\.conv1\.weight', lambda: tf_params[f'synthesis/4x4/Conv/weight'].transpose(3, 2, 0, 1),
r'synthesis\.b4\.conv1\.bias', lambda: tf_params[f'synthesis/4x4/Conv/bias'],
r'synthesis\.b4\.conv1\.noise_const', lambda: tf_params[f'synthesis/noise0'][0, 0],
r'synthesis\.b4\.conv1\.noise_strength', lambda: tf_params[f'synthesis/4x4/Conv/noise_strength'],
r'synthesis\.b4\.conv1\.affine\.weight', lambda: tf_params[f'synthesis/4x4/Conv/mod_weight'].transpose(),
r'synthesis\.b4\.conv1\.affine\.bias', lambda: tf_params[f'synthesis/4x4/Conv/mod_bias'] + 1,
r'synthesis\.b(\d+)\.conv0\.weight', lambda r: tf_params[f'synthesis/{r}x{r}/Conv0_up/weight'][::-1, ::-1].transpose(3, 2, 0, 1),
r'synthesis\.b(\d+)\.conv0\.bias', lambda r: tf_params[f'synthesis/{r}x{r}/Conv0_up/bias'],
r'synthesis\.b(\d+)\.conv0\.noise_const', lambda r: tf_params[f'synthesis/noise{int(np.log2(int(r)))*2-5}'][0, 0],
r'synthesis\.b(\d+)\.conv0\.noise_strength', lambda r: tf_params[f'synthesis/{r}x{r}/Conv0_up/noise_strength'],
r'synthesis\.b(\d+)\.conv0\.affine\.weight', lambda r: tf_params[f'synthesis/{r}x{r}/Conv0_up/mod_weight'].transpose(),
r'synthesis\.b(\d+)\.conv0\.affine\.bias', lambda r: tf_params[f'synthesis/{r}x{r}/Conv0_up/mod_bias'] + 1,
r'synthesis\.b(\d+)\.conv1\.weight', lambda r: tf_params[f'synthesis/{r}x{r}/Conv1/weight'].transpose(3, 2, 0, 1),
r'synthesis\.b(\d+)\.conv1\.bias', lambda r: tf_params[f'synthesis/{r}x{r}/Conv1/bias'],
r'synthesis\.b(\d+)\.conv1\.noise_const', lambda r: tf_params[f'synthesis/noise{int(np.log2(int(r)))*2-4}'][0, 0],
r'synthesis\.b(\d+)\.conv1\.noise_strength', lambda r: tf_params[f'synthesis/{r}x{r}/Conv1/noise_strength'],
r'synthesis\.b(\d+)\.conv1\.affine\.weight', lambda r: tf_params[f'synthesis/{r}x{r}/Conv1/mod_weight'].transpose(),
r'synthesis\.b(\d+)\.conv1\.affine\.bias', lambda r: tf_params[f'synthesis/{r}x{r}/Conv1/mod_bias'] + 1,
r'synthesis\.b(\d+)\.torgb\.weight', lambda r: tf_params[f'synthesis/{r}x{r}/ToRGB/weight'].transpose(3, 2, 0, 1),
r'synthesis\.b(\d+)\.torgb\.bias', lambda r: tf_params[f'synthesis/{r}x{r}/ToRGB/bias'],
r'synthesis\.b(\d+)\.torgb\.affine\.weight', lambda r: tf_params[f'synthesis/{r}x{r}/ToRGB/mod_weight'].transpose(),
r'synthesis\.b(\d+)\.torgb\.affine\.bias', lambda r: tf_params[f'synthesis/{r}x{r}/ToRGB/mod_bias'] + 1,
r'synthesis\.b(\d+)\.skip\.weight', lambda r: tf_params[f'synthesis/{r}x{r}/Skip/weight'][::-1, ::-1].transpose(3, 2, 0, 1),
r'.*\.resample_filter', None,
r'.*\.act_filter', None,
)
return G
#----------------------------------------------------------------------------
def convert_tf_discriminator(tf_D):
if tf_D.version < 4:
raise ValueError('TensorFlow pickle version too low')
# Collect kwargs.
tf_kwargs = tf_D.static_kwargs
known_kwargs = set()
def kwarg(tf_name, default=None):
known_kwargs.add(tf_name)
return tf_kwargs.get(tf_name, default)
# Convert kwargs.
kwargs = dnnlib.EasyDict(
c_dim = kwarg('label_size', 0),
img_resolution = kwarg('resolution', 1024),
img_channels = kwarg('num_channels', 3),
architecture = kwarg('architecture', 'resnet'),
channel_base = kwarg('fmap_base', 16384) * 2,
channel_max = kwarg('fmap_max', 512),
num_fp16_res = kwarg('num_fp16_res', 0),
conv_clamp = kwarg('conv_clamp', None),
cmap_dim = kwarg('mapping_fmaps', None),
block_kwargs = dnnlib.EasyDict(
activation = kwarg('nonlinearity', 'lrelu'),
resample_filter = kwarg('resample_kernel', [1,3,3,1]),
freeze_layers = kwarg('freeze_layers', 0),
),
mapping_kwargs = dnnlib.EasyDict(
num_layers = kwarg('mapping_layers', 0),
embed_features = kwarg('mapping_fmaps', None),
layer_features = kwarg('mapping_fmaps', None),
activation = kwarg('nonlinearity', 'lrelu'),
lr_multiplier = kwarg('mapping_lrmul', 0.1),
),
epilogue_kwargs = dnnlib.EasyDict(
mbstd_group_size = kwarg('mbstd_group_size', None),
mbstd_num_channels = kwarg('mbstd_num_features', 1),
activation = kwarg('nonlinearity', 'lrelu'),
),
)
# Check for unknown kwargs.
kwarg('structure')
kwarg('conditioning')
unknown_kwargs = list(set(tf_kwargs.keys()) - known_kwargs)
if len(unknown_kwargs) > 0:
raise ValueError('Unknown TensorFlow kwarg', unknown_kwargs[0])
# Collect params.
tf_params = _collect_tf_params(tf_D)
for name, value in list(tf_params.items()):
match = re.fullmatch(r'FromRGB_lod(\d+)/(.*)', name)
if match:
r = kwargs.img_resolution // (2 ** int(match.group(1)))
tf_params[f'{r}x{r}/FromRGB/{match.group(2)}'] = value
kwargs.architecture = 'orig'
#for name, value in tf_params.items(): print(f'{name:<50s}{list(value.shape)}')
# Convert params.
from training import networks_stylegan2
D = networks_stylegan2.Discriminator(**kwargs).eval().requires_grad_(False)
# pylint: disable=unnecessary-lambda
# pylint: disable=f-string-without-interpolation
_populate_module_params(D,
r'b(\d+)\.fromrgb\.weight', lambda r: tf_params[f'{r}x{r}/FromRGB/weight'].transpose(3, 2, 0, 1),
r'b(\d+)\.fromrgb\.bias', lambda r: tf_params[f'{r}x{r}/FromRGB/bias'],
r'b(\d+)\.conv(\d+)\.weight', lambda r, i: tf_params[f'{r}x{r}/Conv{i}{["","_down"][int(i)]}/weight'].transpose(3, 2, 0, 1),
r'b(\d+)\.conv(\d+)\.bias', lambda r, i: tf_params[f'{r}x{r}/Conv{i}{["","_down"][int(i)]}/bias'],
r'b(\d+)\.skip\.weight', lambda r: tf_params[f'{r}x{r}/Skip/weight'].transpose(3, 2, 0, 1),
r'mapping\.embed\.weight', lambda: tf_params[f'LabelEmbed/weight'].transpose(),
r'mapping\.embed\.bias', lambda: tf_params[f'LabelEmbed/bias'],
r'mapping\.fc(\d+)\.weight', lambda i: tf_params[f'Mapping{i}/weight'].transpose(),
r'mapping\.fc(\d+)\.bias', lambda i: tf_params[f'Mapping{i}/bias'],
r'b4\.conv\.weight', lambda: tf_params[f'4x4/Conv/weight'].transpose(3, 2, 0, 1),
r'b4\.conv\.bias', lambda: tf_params[f'4x4/Conv/bias'],
r'b4\.fc\.weight', lambda: tf_params[f'4x4/Dense0/weight'].transpose(),
r'b4\.fc\.bias', lambda: tf_params[f'4x4/Dense0/bias'],
r'b4\.out\.weight', lambda: tf_params[f'Output/weight'].transpose(),
r'b4\.out\.bias', lambda: tf_params[f'Output/bias'],
r'.*\.resample_filter', None,
)
return D
#----------------------------------------------------------------------------
@click.command()
@click.option('--source', help='Input pickle', required=True, metavar='PATH')
@click.option('--dest', help='Output pickle', required=True, metavar='PATH')
@click.option('--force-fp16', help='Force the networks to use FP16', type=bool, default=False, metavar='BOOL', show_default=True)
def convert_network_pickle(source, dest, force_fp16):
"""Convert legacy network pickle into the native PyTorch format.
The tool is able to load the main network configurations exported using the TensorFlow version of StyleGAN2 or StyleGAN2-ADA.
It does not support e.g. StyleGAN2-ADA comparison methods, StyleGAN2 configs A-D, or StyleGAN1 networks.
Example:
\b
python legacy.py \\
--source=https://nvlabs-fi-cdn.nvidia.com/stylegan2/networks/stylegan2-cat-config-f.pkl \\
--dest=stylegan2-cat-config-f.pkl
"""
print(f'Loading "{source}"...')
with dnnlib.util.open_url(source) as f:
data = load_network_pkl(f, force_fp16=force_fp16)
print(f'Saving "{dest}"...')
with open(dest, 'wb') as f:
pickle.dump(data, f)
print('Done.')
#----------------------------------------------------------------------------
if __name__ == "__main__":
convert_network_pickle() # pylint: disable=no-value-for-parameter
#----------------------------------------------------------------------------
================================================
FILE: requirements.txt
================================================
torch>=2.0.0
scipy>=1.11.1
Ninja==1.10.2
gradio>=3.35.2
imageio-ffmpeg>=0.4.3
huggingface_hub
hf_transfer
pyopengl
imgui
glfw==2.6.1
pillow>=9.4.0
torchvision>=0.15.2
imageio>=2.9.0
================================================
FILE: stylegan_human/.gitignore
================================================
.DS_Store
__pycache__
*.pt
*.pth
*.pdparams
*.pdiparams
*.pdmodel
*.pkl
*.info
*.yaml
================================================
FILE: stylegan_human/PP_HumanSeg/deploy/infer.py
================================================
# Copyright (c) SenseTime Research. All rights reserved.
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import codecs
import os
import time
import yaml
import numpy as np
import cv2
import paddle
import paddleseg.transforms as T
from paddle.inference import create_predictor, PrecisionType
from paddle.inference import Config as PredictConfig
from paddleseg.core.infer import reverse_transform
from paddleseg.cvlibs import manager
from paddleseg.utils import TimeAverager
from ..scripts.optic_flow_process import optic_flow_process
class DeployConfig:
def __init__(self, path):
with codecs.open(path, 'r', 'utf-8') as file:
self.dic = yaml.load(file, Loader=yaml.FullLoader)
self._transforms = self._load_transforms(self.dic['Deploy'][
'transforms'])
self._dir = os.path.dirname(path)
@property
def transforms(self):
return self._transforms
@property
def model(self):
return os.path.join(self._dir, self.dic['Deploy']['model'])
@property
def params(self):
return os.path.join(self._dir, self.dic['Deploy']['params'])
def _load_transforms(self, t_list):
com = manager.TRANSFORMS
transforms = []
for t in t_list:
ctype = t.pop('type')
transforms.append(com[ctype](**t))
return transforms
class Predictor:
def __init__(self, args):
self.cfg = DeployConfig(args.cfg)
self.args = args
self.compose = T.Compose(self.cfg.transforms)
resize_h, resize_w = args.input_shape
self.disflow = cv2.DISOpticalFlow_create(
cv2.DISOPTICAL_FLOW_PRESET_ULTRAFAST)
self.prev_gray = np.zeros((resize_h, resize_w), np.uint8)
self.prev_cfd = np.zeros((resize_h, resize_w), np.float32)
self.is_init = True
pred_cfg = PredictConfig(self.cfg.model, self.cfg.params)
pred_cfg.disable_glog_info()
if self.args.use_gpu:
pred_cfg.enable_use_gpu(100, 0)
self.predictor = create_predictor(pred_cfg)
if self.args.test_speed:
self.cost_averager = TimeAverager()
def preprocess(self, img):
ori_shapes = []
processed_imgs = []
processed_img = self.compose(img)[0]
processed_imgs.append(processed_img)
ori_shapes.append(img.shape)
return processed_imgs, ori_shapes
def run(self, img, bg):
input_names = self.predictor.get_input_names()
input_handle = self.predictor.get_input_handle(input_names[0])
processed_imgs, ori_shapes = self.preprocess(img)
data = np.array(processed_imgs)
input_handle.reshape(data.shape)
input_handle.copy_from_cpu(data)
if self.args.test_speed:
start = time.time()
self.predictor.run()
if self.args.test_speed:
self.cost_averager.record(time.time() - start)
output_names = self.predictor.get_output_names()
output_handle = self.predictor.get_output_handle(output_names[0])
output = output_handle.copy_to_cpu()
return self.postprocess(output, img, ori_shapes[0], bg)
def postprocess(self, pred, img, ori_shape, bg):
if not os.path.exists(self.args.save_dir):
os.makedirs(self.args.save_dir)
resize_w = pred.shape[-1]
resize_h = pred.shape[-2]
if self.args.soft_predict:
if self.args.use_optic_flow:
score_map = pred[:, 1, :, :].squeeze(0)
score_map = 255 * score_map
cur_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
cur_gray = cv2.resize(cur_gray, (resize_w, resize_h))
optflow_map = optic_flow_process(cur_gray, score_map, self.prev_gray, self.prev_cfd, \
self.disflow, self.is_init)
self.prev_gray = cur_gray.copy()
self.prev_cfd = optflow_map.copy()
self.is_init = False
score_map = np.repeat(optflow_map[:, :, np.newaxis], 3, axis=2)
score_map = np.transpose(score_map, [2, 0, 1])[np.newaxis, ...]
score_map = reverse_transform(
paddle.to_tensor(score_map),
ori_shape,
self.cfg.transforms,
mode='bilinear')
alpha = np.transpose(score_map.numpy().squeeze(0),
[1, 2, 0]) / 255
else:
score_map = pred[:, 1, :, :]
score_map = score_map[np.newaxis, ...]
score_map = reverse_transform(
paddle.to_tensor(score_map),
ori_shape,
self.cfg.transforms,
mode='bilinear')
alpha = np.transpose(score_map.numpy().squeeze(0), [1, 2, 0])
else:
if pred.ndim == 3:
pred = pred[:, np.newaxis, ...]
result = reverse_transform(
paddle.to_tensor(
pred, dtype='float32'),
ori_shape,
self.cfg.transforms,
mode='bilinear')
result = np.array(result)
if self.args.add_argmax:
result = np.argmax(result, axis=1)
else:
result = result.squeeze(1)
alpha = np.transpose(result, [1, 2, 0])
# background replace
h, w, _ = img.shape
if bg is None:
bg = np.ones_like(img)*255
else:
bg = cv2.resize(bg, (w, h))
if bg.ndim == 2:
bg = bg[..., np.newaxis]
comb = (alpha * img + (1 - alpha) * bg).astype(np.uint8)
return comb, alpha, bg, img
================================================
FILE: stylegan_human/PP_HumanSeg/export_model/download_export_model.py
================================================
# coding: utf8
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import sys
import os
LOCAL_PATH = os.path.dirname(os.path.abspath(__file__))
TEST_PATH = os.path.join(LOCAL_PATH, "../../../", "test")
sys.path.append(TEST_PATH)
from paddleseg.utils.download import download_file_and_uncompress
model_urls = {
"pphumanseg_lite_portrait_398x224_with_softmax":
"https://paddleseg.bj.bcebos.com/dygraph/ppseg/ppseg_lite_portrait_398x224_with_softmax.tar.gz",
"deeplabv3p_resnet50_os8_humanseg_512x512_100k_with_softmax":
"https://paddleseg.bj.bcebos.com/dygraph/humanseg/export/deeplabv3p_resnet50_os8_humanseg_512x512_100k_with_softmax.zip",
"fcn_hrnetw18_small_v1_humanseg_192x192_with_softmax":
"https://paddleseg.bj.bcebos.com/dygraph/humanseg/export/fcn_hrnetw18_small_v1_humanseg_192x192_with_softmax.zip",
"pphumanseg_lite_generic_humanseg_192x192_with_softmax":
"https://paddleseg.bj.bcebos.com/dygraph/humanseg/export/pphumanseg_lite_generic_192x192_with_softmax.zip",
}
if __name__ == "__main__":
for model_name, url in model_urls.items():
download_file_and_uncompress(
url=url,
savepath=LOCAL_PATH,
extrapath=LOCAL_PATH,
extraname=model_name)
print("Export model download success!")
================================================
FILE: stylegan_human/PP_HumanSeg/pretrained_model/download_pretrained_model.py
================================================
# coding: utf8
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import sys
import os
LOCAL_PATH = os.path.dirname(os.path.abspath(__file__))
TEST_PATH = os.path.join(LOCAL_PATH, "../../../", "test")
sys.path.append(TEST_PATH)
from paddleseg.utils.download import download_file_and_uncompress
model_urls = {
"pphumanseg_lite_portrait_398x224":
"https://paddleseg.bj.bcebos.com/dygraph/ppseg/ppseg_lite_portrait_398x224.tar.gz",
"deeplabv3p_resnet50_os8_humanseg_512x512_100k":
"https://paddleseg.bj.bcebos.com/dygraph/humanseg/train/deeplabv3p_resnet50_os8_humanseg_512x512_100k.zip",
"fcn_hrnetw18_small_v1_humanseg_192x192":
"https://paddleseg.bj.bcebos.com/dygraph/humanseg/train/fcn_hrnetw18_small_v1_humanseg_192x192.zip",
"pphumanseg_lite_generic_human_192x192":
"https://paddleseg.bj.bcebos.com/dygraph/humanseg/train/pphumanseg_lite_generic_192x192.zip",
}
if __name__ == "__main__":
for model_name, url in model_urls.items():
download_file_and_uncompress(
url=url,
savepath=LOCAL_PATH,
extrapath=LOCAL_PATH,
extraname=model_name)
print("Pretrained model download success!")
================================================
FILE: stylegan_human/README.md
================================================
# StyleGAN-Human: A Data-Centric Odyssey of Human Generation
<img src="./img/demo_V5_thumbnails-min.png" width="96%" height="96%">
<!--
**stylegan-human/StyleGAN-Human** is a ✨ _special_ ✨ repository because its `README.md` (this file) appears on your GitHub profile.
-->
>
>
> **Abstract:** *Unconditional human image generation is an important task in vision and graphics, which enables various applications in the creative industry. Existing studies in this field mainly focus on "network engineering" such as designing new components and objective functions. This work takes a data-centric perspective and investigates multiple critical aspects in "data engineering", which we believe would complement the current practice. To facilitate a comprehensive study, we collect and annotate a large-scale human image dataset with over 230K samples capturing diverse poses and textures. Equipped with this large dataset, we rigorously investigate three essential factors in data engineering for StyleGAN-based human generation, namely data size, data distribution, and data alignment. Extensive experiments reveal several valuable observations w.r.t. these aspects: 1) Large-scale data, more than 40K images, are needed to train a high-fidelity unconditional human generation model with vanilla StyleGAN. 2) A balanced training set helps improve the generation quality with rare face poses compared to the long-tailed counterpart, whereas simply balancing the clothing texture distribution does not effectively bring an improvement. 3) Human GAN models with body centers for alignment outperform models trained using face centers or pelvis points as alignment anchors. In addition, a model zoo and human editing applications are demonstrated to facilitate future research in the community.* <br>
**Keyword:** Human Image Generation, Data-Centric, StyleGAN
[Jianglin Fu](mailto:fujianglin@sensetime.com), [Shikai Li](mailto:lishikai@sensetime.com), [Yuming Jiang](https://yumingj.github.io/), [Kwan-Yee Lin](https://kwanyeelin.github.io/), [Chen Qian](https://scholar.google.com/citations?user=AerkT0YAAAAJ&hl=zh-CN), [Chen Change Loy](https://www.mmlab-ntu.com/person/ccloy/), [Wayne Wu](https://wywu.github.io/), and [Ziwei Liu](https://liuziwei7.github.io/) <br>
**[[Demo Video]](https://youtu.be/nIrb9hwsdcI)** | **[[Project Page]](https://stylegan-human.github.io/)** | **[[Paper]](https://arxiv.org/pdf/2204.11823.pdf)**
## Updates
- [20/07/2022] [SHHQ-1.0](./docs/Dataset.md) dataset with 40K images is released! :sparkles:
- [15/06/2022] Data alignment and real-image inversion scripts are released.
- [26/04/2022] Technical report released!
- [22/04/2022] Technical report will be released before May.
- [21/04/2022] The codebase and project page are created.
## Data Download
The first version SHHQ-1.0, with 40K images is released. To download and use the dataset set, please read the instructions in [Dataset.md](./docs/Dataset.md)
（We are currently facing large incoming applications, and we need to carefully verify all the applicants, please be patient, and we will reply to you as soon as possible.）
## Model Zoo
| Structure | 1024x512 | Metric | Scores | 512x256 | Metric | Scores |
| --------- |:----------:| :----------:| :----------:| :-----: | :-----: | :-----: |
| StyleGAN1 |[stylegan_human_v1_1024.pkl](https://drive.google.com/file/d/1h-R-IV-INGdPEzj4P9ml6JTEvihuNgLX/view?usp=sharing)| fid50k | 3.79 | to be released | - | - |
| StyleGAN2 |[stylegan_human_v2_1024.pkl](https://drive.google.com/file/d/1FlAb1rYa0r_--Zj_ML8e6shmaF28hQb5/view?usp=sharing)| fid50k_full | 1.57 |[stylegan_human_v2_512.pkl](https://drive.google.com/file/d/1dlFEHbu-WzQWJl7nBBZYcTyo000H9hVm/view?usp=sharing) | fid50k_full | 1.97 |
| StyleGAN3 |to be released | - | - | [stylegan_human_v3_512.pkl](https://drive.google.com/file/d/1_274jk_N6WSCkKWeu7hjHycqGvbuOFf5/view?usp=sharing) | fid50k_full | 2.54 |
## Web Demo
Integrated into [Huggingface Spaces 🤗](https://huggingface.co/spaces) using [Gradio](https://github.com/gradio-app/gradio). Try out the Web Demo for generation: [](https://huggingface.co/spaces/hysts/StyleGAN-Human) and interpolation [](https://huggingface.co/spaces/hysts/StyleGAN-Human-Interpolation)
<a href="https://colab.research.google.com/drive/1sgxoDM55iM07FS54vz9ALg1XckiYA2On"><img src="https://colab.research.google.com/assets/colab-badge.svg" height=22.5></a>
We prepare a Colab demo to allow you to synthesize images with the provided models, as well as visualize the performance of style-mixing, interpolation, and attributes editing.
The notebook will guide you to install the necessary environment and download pretrained models. The output images can be found in `./StyleGAN-Human/outputs/`.
Hope you enjoy!
## Usage
### System requirements
* The original code bases are [stylegan (tensorflow)](https://github.com/NVlabs/stylegan), [stylegan2-ada (pytorch)](https://github.com/NVlabs/stylegan2-ada-pytorch), [stylegan3 (pytorch)](https://github.com/NVlabs/stylegan3), released by NVidia
* We tested in Python 3.8.5 and PyTorch 1.9.1 with CUDA 11.1. (See https://pytorch.org for PyTorch install instructions.)
### Installation
To work with this project on your own machine, you need to install the environmnet as follows:
```
conda env create -f environment.yml
conda activate stylehuman
# [Optional: tensorflow 1.x is required for StyleGAN1. ]
pip install nvidia-pyindex
pip install nvidia-tensorflow[horovod]
pip install nvidia-tensorboard==1.15
```
Extra notes:
1. In case having some conflicts when calling CUDA version, please try to empty the LD_LIBRARY_PATH. For example:
```
LD_LIBRARY_PATH=; python generate.py --outdir=out/stylegan_human_v2_1024 --trunc=1 --seeds=1,3,5,7
--network=pretrained_models/stylegan_human_v2_1024.pkl --version 2
```
2. We found the following troubleshooting links might be helpful: [1.](https://github.com/NVlabs/stylegan3), [2.](https://github.com/NVlabs/stylegan3/blob/main/docs/troubleshooting.md)
### Train
The training scripts are based on the original [stylegan1](https://github.com/NVlabs/stylegan), [stylegan2-ada](https://github.com/NVlabs/stylegan2-ada-pytorch), and [stylegan3](https://github.com/NVlabs/stylegan3) with minor changes. Here we only provide the scripts with modifications for SG2 and SG3. You can replace the old files with the provided scripts to train. (assume SHHQ-1.0 is placed under data/)
#### Train Stylegan2-ada-pytorch with SHHQ-1.0
```
python train.py --outdir=training_results/sg2/ --data=data/SHHQ-1.0/ \
--gpus=8 --aug=noaug --mirror=1 --snap=250 --cfg=shhq --square=False
```
#### Train Stylegan3 with SHHQ-1.0
```
python train.py --outdir=training_results/sg3/ --cfg=stylegan3-r --gpus=8 --batch=32 --gamma=12.4 \
--mirror=1 --aug=noaug --data=data/SHHQ-1.0/ --square=False --snap=250
```
### Pretrained models
Please put the downloaded pretrained models [from above link](#Model-Zoo) under the folder 'pretrained_models'.
### Generate full-body human images using our pretrained model
```
# Generate human full-body images without truncation
python generate.py --outdir=outputs/generate/stylegan_human_v2_1024 --trunc=1 --seeds=1,3,5,7 --network=pretrained_models/stylegan_human_v2_1024.pkl --version 2
# Generate human full-body images with truncation
python generate.py --outdir=outputs/generate/stylegan_human_v2_1024 --trunc=0.8 --seeds=0-10 --network=pretrained_models/stylegan_human_v2_1024.pkl --version 2
# Generate human full-body images using stylegan V1
python generate.py --outdir=outputs/generate/stylegan_human_v1_1024 --network=pretrained_models/stylegan_human_v1_1024.pkl --version 1 --seeds=1,3,5
# Generate human full-body images using stylegan V3
python generate.py --outdir=outputs/generate/stylegan_human_v3_512 --network=pretrained_models/stylegan_human_v3_512.pkl --version 3 --seeds=1,3,5
```
#### Note: The following demos are generated based on models related to StyleGAN V2 (stylegan_human_v2_512.pkl and stylegan_human_v2_1024.pkl). If you want to see results for V1 or V3, you need to change the loading method of the corresponding models.
### Interpolation
```
python interpolation.py --network=pretrained_models/stylegan_human_v2_1024.pkl --seeds=85,100 --outdir=outputs/inter_gifs
```
### Style-mixing **image** using stylegan2
```
python style_mixing.py --network=pretrained_models/stylegan_human_v2_1024.pkl --rows=85,100,75,458,1500 \\
--cols=55,821,1789,293 --styles=0-3 --outdir=outputs/stylemixing
```
### Style-mixing **video** using stylegan2
```
python stylemixing_video.py --network=pretrained_models/stylegan_human_v2_1024.pkl --row-seed=3859 \\
--col-seeds=3098,31759,3791 --col-styles=8-12 --trunc=0.8 --outdir=outputs/stylemixing_video
```
### Aligned raw images
For alignment, we use [openpose-pytorch](https://github.com/Hzzone/pytorch-openpose) for body-keypoints detection and [PaddlePaddle](https://github.com/PaddlePaddle/PaddleSeg/tree/release/2.5/contrib/PP-HumanSeg) for human segmentation.
Before running the alignment script, few models need to be installed:
1. download [body_pose_model.pth](https://drive.google.com/drive/folders/1JsvI4M4ZTg98fmnCZLFM-3TeovnCRElG?usp=sharing) and place it into openpose/model/.
2. download and extract [deeplabv3p_resnet50_os8_humanseg_512x512_100k_with_softmax](https://paddleseg.bj.bcebos.com/dygraph/humanseg/export/deeplabv3p_resnet50_os8_humanseg_512x512_100k_with_softmax.zip) into PP_HumanSeg/export_model/deeplabv3p_resnet50_os8_humanseg_512x512_100k_with_softmax.
3. download and extract [deeplabv3p_resnet50_os8_humanseg_512x512_100k](https://paddleseg.bj.bcebos.com/dygraph/humanseg/train/deeplabv3p_resnet50_os8_humanseg_512x512_100k.zip) into PP_HumanSeg/pretrained_model/deeplabv3p_resnet50_os8_humanseg_512x512_100k.
4. install paddlepaddel: ``` pip install paddleseg ```
Then you can start alignment:
```
python alignment.py --image-folder img/test/ --output-folder aligned_image/
```
### Invert real image with [PTI](https://github.com/danielroich/PTI)
Before inversion, please download our PTI weights: [e4e_w+.pt](https://drive.google.com/file/d/1NUfSJqLhsrU7c9PwAtlZ9xtrxhzS_6tu/view?usp=sharing) into /pti/.
Few parameters you can change:
- /pti/pti_configs/hyperparameters.py:
- first_inv_type = 'w+' -> Use pretrained e4e encoder
- first_inv_type = 'w' -> Use projection and optimization
- /pti/pti_configs/paths_config.py:
- input_data_path: path of real images
- e4e: path of e4e_w+.pt
- stylegan2_ada_shhq: pretrained stylegan2-ada model for SHHQ
```
python run_pti.py
```
Note: we used the test image under 'aligned_image/' (the output of alignment.py), the inverted latent code and fine-tuned generator will be saved in 'outputs/pti/'
### Editing with InterfaceGAN, StyleSpace, and Sefa
```
python edit.py --network pretrained_models/stylegan_human_v2_1024.pkl --attr_name upper_length \\
--seeds 61531,61570,61571,61610 --outdir outputs/edit_results
```
### Editing using inverted latent code
```
python edit.py ---network outputs/pti/checkpoints/model_test.pkl --attr_name upper_length \\
--outdir outputs/edit_results --real True --real_w_path outputs/pti/embeddings/test/PTI/test/0.pt --real_img_path aligned_image/test.png
```
Note:
1. ''upper_length'' and ''bottom_length'' of ''attr_name'' are available for demo.
2. Layers to control and editing strength are set in edit/edit_config.py.
### Demo for [InsetGAN](https://arxiv.org/abs/2203.07293)
We implement a quick demo using the key idea from InsetGAN: combining the face generated by FFHQ with the human-body generated by our pretrained model, optimizing both face and body latent codes to get a coherent full-body image.
Before running the script, you need to download the [FFHQ face model]( https://docs.google.com/uc?export=download&confirm=t&id=125OG7SMkXI-Kf2aqiwLLHyCvSW-gZk3M), or you can use your own face model, as well as [pretrained face landmark](https://docs.google.com/uc?export=download&confirm=&id=1A82DnJBJzt8wI2J8ZrCK5fgHcQ2-tcWM) and [pretrained CNN face detection model for dlib](https://docs.google.com/uc?export=download&confirm=&id=1MduBgju5KFNrQfDLoQXJ_1_h5MnctCIG)
```
python insetgan.py --body_network=pretrained_models/stylegan_human_v2_1024.pkl --face_network=pretrained_models/ffhq.pkl \\
--body_seed=82 --face_seed=43 --trunc=0.6 --outdir=outputs/insetgan/ --video 1
```
## Results
### Editing with inverted real image
(from left to right: real image | inverted image | InterFaceGAN result | StyleSpace result | SeFa result)
https://user-images.githubusercontent.com/98547009/173773800-bb7fe54a-84d3-4b30-9864-a6b7b311f8ff.mp4
### For more demo, please visit our [**web page**](https://stylegan-human.github.io/) .
## TODO List
- [ ] Release 1024x512 version of StyleGAN-Human based on StyleGAN3
- [ ] Release 512x256 version of StyleGAN-Human based on StyleGAN1
- [ ] Extension of downstream application (InsetGAN): Add face inversion interface to support fusing user face image and stylegen-human body image
- [x] Add Inversion Script into the provided editing pipeline
- [ ] Release Dataset
## Related Works
* (SIGGRAPH 2022) **Text2Human: Text-Driven Controllable Human Image Generation**, Yuming Jiang et al. [[Paper](https://arxiv.org/pdf/2205.15996.pdf)], [[Code](https://github.com/yumingj/Text2Human)], [[Project Page](https://yumingj.github.io/projects/Text2Human.html)], [[Dataset](https://github.com/yumingj/DeepFashion-MultiModal)]
* (ICCV 2021) **Talk-to-Edit: Fine-Grained Facial Editing via Dialog**, Yuming Jiang et al. [[Paper](https://arxiv.org/abs/2109.04425)], [[Code](https://github.com/yumingj/Talk-to-Edit)], [[Project Page](https://www.mmlab-ntu.com/project/talkedit/)], [[Dataset](https://mmlab.ie.cuhk.edu.hk/projects/CelebA/CelebA_Dialog.html)]
* (Technical Report 2022) **Generalizable Neural Performer: Learning Robust Radiance Fields for Human Novel View Synthesis**, Wei Cheng et al. [[Paper](https://arxiv.org/pdf/2204.11798.pdf)], [[Code](https://github.com/generalizable-neural-performer/gnr)], [[Project Page](https://generalizable-neural-performer.github.io/)], [[Dataset](https://generalizable-neural-performer.github.io/genebody.html)]
## Citation
If you find this work useful for your research, please consider citing our paper:
```bibtex
@article{fu2022styleganhuman,
title={StyleGAN-Human: A Data-Centric Odyssey of Human Generation},
author={Fu, Jianglin and Li, Shikai and Jiang, Yuming and Lin, Kwan-Yee and Qian, Chen and Loy, Chen-Change and Wu, Wayne and Liu, Ziwei},
journal = {arXiv preprint},
volume = {arXiv:2204.11823},
year = {2022}
```
## Acknowlegement
Part of the code is borrowed from [stylegan (tensorflow)](https://github.com/NVlabs/stylegan), [stylegan2-ada (pytorch)](https://github.com/NVlabs/stylegan2-ada-pytorch), [stylegan3 (pytorch)](https://github.com/NVlabs/stylegan3).
================================================
FILE: stylegan_human/__init__.py
================================================
================================================
FILE: stylegan_human/alignment.py
================================================
# Copyright (c) SenseTime Research. All rights reserved.
import os
import argparse
import numpy as np
import torch
from torch.utils.data import DataLoader
from torchvision.transforms import transforms
from utils.ImagesDataset import ImagesDataset
import cv2
import time
import copy
import imutils
# for openpose body keypoint detector : # (src:https://github.com/Hzzone/pytorch-openpose)
from openpose.src import util
from openpose.src.body import Body
# for paddlepaddle human segmentation : #(src: https://github.com/PaddlePaddle/PaddleSeg/blob/release/2.5/contrib/PP-HumanSeg/)
from PP_HumanSeg.deploy.infer import Predictor as PP_HumenSeg_Predictor
import math
def angle_between_points(p0,p1,p2):
if p0[1]==-1 or p1[1]==-1 or p2[1]==-1:
return -1
a = (p1[0]-p0[0])**2 + (p1[1]-p0[1])**2
b = (p1[0]-p2[0])**2 + (p1[1]-p2[1])**2
c = (p2[0]-p0[0])**2 + (p2[1]-p0[1])**2
if a * b == 0:
return -1
return math.acos((a+b-c) / math.sqrt(4*a*b)) * 180 / math.pi
def crop_img_with_padding(img, keypoints, rect):
person_xmin,person_xmax, ymin, ymax= rect
img_h,img_w,_ = img.shape ## find body center using keypoints
middle_shoulder_x = keypoints[1][0]
middle_hip_x = (keypoints[8][0] + keypoints[11][0]) // 2
mid_x = (middle_hip_x + middle_shoulder_x) // 2
mid_y = (ymin + ymax) // 2
## find which side (l or r) is further than center x, use the further side
if abs(mid_x-person_xmin) > abs(person_xmax-mid_x): #left further
xmin = person_xmin
xmax = mid_x + (mid_x-person_xmin)
else:
############### may be negtive
### in this case, the script won't output any image, leave the case like this
### since we don't want to pad human body
xmin = mid_x - (person_xmax-mid_x)
xmax = person_xmax
w = xmax - xmin
h = ymax - ymin
## pad rectangle to w:h = 1:2 ## calculate desired border length
if h / w >= 2: #pad horizontally
target_w = h // 2
xmin_prime = int(mid_x - target_w / 2)
xmax_prime = int(mid_x + target_w / 2)
if xmin_prime < 0:
pad_left = abs(xmin_prime)# - xmin
xmin = 0
else:
pad_left = 0
xmin = xmin_prime
if xmax_prime > img_w:
pad_right = xmax_prime - img_w
xmax = img_w
else:
pad_right = 0
xmax = xmax_prime
cropped_img = img[int(ymin):int(ymax), int(xmin):int(xmax)]
im_pad = cv2.copyMakeBorder(cropped_img, 0, 0, int(pad_left), int(pad_right), cv2.BORDER_REPLICATE)
else: #pad vertically
target_h = w * 2
ymin_prime = mid_y - (target_h / 2)
ymax_prime = mid_y + (target_h / 2)
if ymin_prime < 0:
pad_up = abs(ymin_prime)# - ymin
ymin = 0
else:
pad_up = 0
ymin = ymin_prime
if ymax_prime > img_h:
pad_down = ymax_prime - img_h
ymax = img_h
else:
pad_down = 0
ymax = ymax_prime
print(ymin,ymax, xmin,xmax, img.shape)
cropped_img = img[int(ymin):int(ymax), int(xmin):int(xmax)]
im_pad = cv2.copyMakeBorder(cropped_img, int(pad_up), int(pad_down), 0,
0, cv2.BORDER_REPLICATE)
result = cv2.resize(im_pad,(512,1024),interpolation = cv2.INTER_AREA)
return result
def run(args):
os.makedirs(args.output_folder, exist_ok=True)
dataset = ImagesDataset(args.image_folder, transforms.Compose([transforms.ToTensor()]))
dataloader = DataLoader(dataset, batch_size=1, shuffle=False)
body_estimation = Body('openpose/model/body_pose_model.pth')
total = len(dataloader)
print('Num of dataloader : ', total)
os.makedirs(f'{args.output_folder}', exist_ok=True)
# os.makedirs(f'{args.output_folder}/middle_result', exist_ok=True)
## initialzide HumenSeg
human_seg_args = {}
human_seg_args['cfg'] = 'PP_HumanSeg/export_model/deeplabv3p_resnet50_os8_humanseg_512x512_100k_with_softmax/deploy.yaml'
human_seg_args['input_shape'] = [1024,512]
human_seg_args['save_dir'] = args.output_folder
human_seg_args['soft_predict'] = False
human_seg_args['use_gpu'] = True
human_seg_args['test_speed'] = False
human_seg_args['use_optic_flow'] = False
human_seg_args['add_argmax'] = True
human_seg_args= argparse.Namespace(**human_seg_args)
human_seg = PP_HumenSeg_Predictor(human_seg_args)
from tqdm import tqdm
for fname, image in tqdm(dataloader):
# try:
## tensor to numpy image
fname = fname[0]
print(f'Processing \'{fname}\'.')
image = (image.permute(0, 2, 3, 1) * 255).clamp(0, 255)
image = image.squeeze(0).numpy() # --> tensor to numpy, (H,W,C)
# avoid super high res img
if image.shape[0] >= 2000: # height ### for shein image
ratio = image.shape[0]/1200 #height
dim = (int(image.shape[1]/ratio),1200)#(width, height)
image = cv2.resize(image, dim, interpolation = cv2.INTER_AREA)
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
## create segmentation
# mybg = cv2.imread('mybg.png')
comb, segmentation, bg, ori_img = human_seg.run(image,None) #mybg)
# cv2.imwrite('comb.png',comb) # [0,255]
# cv2.imwrite('alpha.png',segmentation*255) # segmentation [0,1] --> [0.255]
# cv2.imwrite('bg.png',bg) #[0,255]
# cv2.imwrite('ori_img.png',ori_img) # [0,255]
masks_np = (segmentation* 255)# .byte().cpu().numpy() #1024,512,1
mask0_np = masks_np[:,:,0].astype(np.uint8)#[0, :, :]
contours = cv2.findContours(mask0_np, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
cnts = imutils.grab_contours(contours)
c = max(cnts, key=cv2.contourArea)
extTop = tuple(c[c[:, :, 1].argmin()][0])
extBot = tuple(c[c[:, :, 1].argmax()][0])
extBot = list(extBot)
extTop = list(extTop)
pad_range = int((extBot[1]-extTop[1])*0.05)
if (int(extTop[1])<=5 and int(extTop[1])>0) and (comb.shape[0]>int(extBot[1]) and int(extBot[1])>=comb.shape[0]-5): #seg mask already reaches to the edge
#pad with pure white, top 100 px, bottom 100 px
comb= cv2.copyMakeBorder(comb,pad_range+5,pad_range+5,0,0,cv2.BORDER_CONSTANT,value=[255,255,255])
elif int(extTop[1])<=0 or int(extBot[1])>=comb.shape[0]:
print('PAD: body out of boundary', fname) #should not happened
return {}
else:
comb = cv2.copyMakeBorder(comb, pad_range+5, pad_range+5, 0, 0, cv2.BORDER_REPLICATE) #105 instead of 100: give some extra space
extBot[1] = extBot[1] + pad_range+5
extTop[1] = extTop[1] + pad_range+5
extLeft = tuple(c[c[:, :, 0].argmin()][0])
extRight = tuple(c[c[:, :, 0].argmax()][0])
extLeft = list(extLeft)
extRight = list(extRight)
person_ymin = int(extTop[1])-pad_range # 100
person_ymax = int(extBot[1])+pad_range # 100 #height
if person_ymin<0 or person_ymax>comb.shape[0]: # out of range
return {}
person_xmin = int(extLeft[0])
person_xmax = int(extRight[0])
rect = [person_xmin,person_xmax,person_ymin, person_ymax]
# recimg = copy.deepcopy(comb)
# cv2.rectangle(recimg,(person_xmin,person_ymin),(person_xmax,person_ymax),(0,255,0),2)
# cv2.imwrite(f'{args.output_folder}/middle_result/{fname}_rec.png',recimg)
## detect keypoints
keypoints, subset = body_estimation(comb)
# print(keypoints, subset, len(subset))
if len(subset) != 1 or (len(subset)==1 and subset[0][-1]<15):
print(f'Processing \'{fname}\'. Please import image contains one person only. Also can check segmentation mask. ')
continue
# canvas = copy.deepcopy(comb)
# canvas = util.draw_bodypose(canvas, keypoints, subset, show_number=True)
# cv2.imwrite(f'{args.output_folder}/middle_result/{fname}_keypoints.png',canvas)
comb = crop_img_with_padding(comb, keypoints, rect)
cv2.imwrite(f'{args.output_folder}/{fname}.png', comb)
print(f' -- Finished processing \'{fname}\'. --')
# except:
# print(f'Processing \'{fname}\'. Not satisfied the alignment strategy.')
if __name__ == '__main__':
torch.backends.cudnn.benchmark = True
torch.backends.cudnn.deterministic = False
t1 = time.time()
arg_formatter = argparse.ArgumentDefaultsHelpFormatter
description = 'StyleGAN-Human data process'
parser = argparse.ArgumentParser(formatter_class=arg_formatter,
description=description)
parser.add_argument('--image-folder', type=str, dest='image_folder')
parser.add_argument('--output-folder', dest='output_folder', default='results', type=str)
# parser.add_argument('--cfg', dest='cfg for segmentation', default='PP_HumanSeg/export_model/ppseg_lite_portrait_398x224_with_softmax/deploy.yaml', type=str)
print('parsing arguments')
cmd_args = parser.parse_args()
run(cmd_args)
print('total time elapsed: ', str(time.time() - t1))
================================================
FILE: stylegan_human/bg_white.py
================================================
# Copyright (c) SenseTime Research. All rights reserved.
import os
import click
import cv2
import numpy as np
def bg_white(seg, raw, blur_level=3, gaussian=81):
seg = cv2.blur(seg, (blur_level, blur_level))
empty = np.ones_like(seg)
seg_bg = (empty - seg) * 255
seg_bg = cv2.GaussianBlur(seg_bg,(gaussian,gaussian),0)
background_mask = cv2.cvtColor(255 - cv2.cvtColor(seg, cv2.COLOR_BGR2GRAY), cv2.COLOR_GRAY2BGR)
masked_fg = (raw * (1 / 255)) * (seg * (1 / 255))
masked_bg = (seg_bg * (1 / 255)) * (background_mask * (1 / 255))
frame = np.uint8(cv2.add(masked_bg,masked_fg)*255)
return frame
"""
To turn background into white.
Examples:
\b
python bg_white.py --raw_img_dir=./SHHQ-1.0/no_segment/ --raw_seg_dir=./SHHQ-1.0/segments/ \\
--outdir=./SHHQ-1.0/bg_white/
"""
@click.command()
@click.pass_context
@click.option('--raw_img_dir', default="./SHHQ-1.0/no_segment/", help='folder of raw image', required=True)
@click.option('--raw_seg_dir', default='./SHHQ-1.0/segments/', help='folder of segmentation masks', required=True)
@click.option('--outdir', help='Where to save the output images', default= "./SHHQ-1.0/bg_white/" , type=str, required=True, metavar='DIR')
def main(
ctx: click.Context,
raw_img_dir: str,
raw_seg_dir: str,
outdir: str):
os.makedirs(outdir, exist_ok=True)
files = os.listdir(raw_img_dir)
for file in files:
print(file)
raw = cv2.imread(os.path.join(raw_img_dir, file))
seg = cv2.imread(os.path.join(raw_seg_dir, file))
assert raw is not None
assert seg is not None
white_frame = bg_white(seg, raw)
cv2.imwrite(os.path.join(outdir,file), white_frame)
if __name__ == "__main__":
main()
================================================
FILE: stylegan_human/dnnlib/__init__.py
================================================
# Copyright (c) SenseTime Research. All rights reserved.
# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
#
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto. Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.
from .util import EasyDict, make_cache_dir_path
================================================
FILE: stylegan_human/dnnlib/tflib/__init__.py
================================================
# Copyright (c) SenseTime Research. All rights reserved.
# Copyright (c) 2019, NVIDIA Corporation. All rights reserved.
#
# This work is made available under the Nvidia Source Code License-NC.
# To view a copy of this license, visit
# https://nvlabs.github.io/stylegan2/license.html
from . import autosummary
from . import network
from . import optimizer
from . import tfutil
from . import custom_ops
from .tfutil import *
from .network import Network
from .optimizer import Optimizer
from .custom_ops import get_plugin
================================================
FILE: stylegan_human/dnnlib/tflib/autosummary.py
================================================
# Copyright (c) SenseTime Research. All rights reserved.
# Copyright (c) 2019, NVIDIA Corporation. All rights reserved.
#
# This work is made available under the Nvidia Source Code License-NC.
# To view a copy of this license, visit
# https://nvlabs.github.io/stylegan2/license.html
"""Helper for adding automatically tracked values to Tensorboard.
Autosummary creates an identity op that internally keeps track of the input
values and automatically shows up in TensorBoard. The reported value
represents an average over input components. The average is accumulated
constantly over time and flushed when save_summaries() is called.
Notes:
- The output tensor must be used as an input for something else in the
graph. Otherwise, the autosummary op will not get executed, and the average
value will not get accumulated.
- It is perfectly fine to include autosummaries with the same name in
several places throughout the graph, even if they are executed concurrently.
- It is ok to also pass in a python scalar or numpy array. In this case, it
is added to the average immediately.
"""
from collections import OrderedDict
import numpy as np
import tensorflow as tf
from tensorboard import summary as summary_lib
from tensorboard.plugins.custom_scalar import layout_pb2
from . import tfutil
from .tfutil import TfExpression
from .tfutil import TfExpressionEx
# Enable "Custom scalars" tab in TensorBoard for advanced formatting.
# Disabled by default to reduce tfevents file size.
enable_custom_scalars = False
_dtype = tf.float64
_vars = OrderedDict() # name => [var, ...]
_immediate = OrderedDict() # name => update_op, update_value
_finalized = False
_merge_op = None
def _create_var(name: str, value_expr: TfExpression) -> TfExpression:
"""Internal helper for creating autosummary accumulators."""
assert not _finalized
name_id = name.replace("/", "_")
v = tf.cast(value_expr, _dtype)
if v.shape.is_fully_defined():
size = np.prod(v.shape.as_list())
size_expr = tf.constant(size, dtype=_dtype)
else:
size = None
size_expr = tf.reduce_prod(tf.cast(tf.shape(v), _dtype))
if size == 1:
if v.shape.ndims != 0:
v = tf.reshape(v, [])
v = [size_expr, v, tf.square(v)]
else:
v = [size_expr, tf.reduce_sum(v), tf.reduce_sum(tf.square(v))]
v = tf.cond(tf.is_finite(v[1]), lambda: tf.stack(v), lambda: tf.zeros(3, dtype=_dtype))
with tfutil.absolute_name_scope("Autosummary/" + name_id), tf.control_dependencies(None):
var = tf.Variable(tf.zeros(3, dtype=_dtype), trainable=False) # [sum(1), sum(x), sum(x**2)]
update_op = tf.cond(tf.is_variable_initialized(var), lambda: tf.assign_add(var, v), lambda: tf.assign(var, v))
if name in _vars:
_vars[name].append(var)
else:
_vars[name] = [var]
return update_op
def autosummary(name: str, value: TfExpressionEx, passthru: TfExpressionEx = None, condition: TfExpressionEx = True) -> TfExpressionEx:
"""Create a new autosummary.
Args:
name: Name to use in TensorBoard
value: TensorFlow expression or python value to track
passthru: Optionally return this TF node without modifications but tack an autosummary update side-effect to this node.
Example use of the passthru mechanism:
n = autosummary('l2loss', loss, passthru=n)
This is a shorthand for the following code:
with tf.control_dependencies([autosummary('l2loss', loss)]):
n = tf.identity(n)
"""
tfutil.assert_tf_initialized()
name_id = name.replace("/", "_")
if tfutil.is_tf_expression(value):
with tf.name_scope("summary_" + name_id), tf.device(value.device):
condition = tf.convert_to_tensor(condition, name='condition')
update_op = tf.cond(condition, lambda: tf.group(_create_var(name, value)), tf.no_op)
with tf.control_dependencies([update_op]):
return tf.identity(value if passthru is None else passthru)
else: # python scalar or numpy array
assert not tfutil.is_tf_expression(passthru)
assert not tfutil.is_tf_expression(condition)
if condition:
if name not in _immediate:
with tfutil.absolute_name_scope("Autosummary/" + name_id), tf.device(None), tf.control_dependencies(None):
update_value = tf.placeholder(_dtype)
update_op = _create_var(name, update_value)
_immediate[name] = update_op, update_value
update_op, update_value = _immediate[name]
tfutil.run(update_op, {update_value: value})
return value if passthru is None else passthru
def finalize_autosummaries() -> None:
"""Create the necessary ops to include autosummaries in TensorBoard report.
Note: This should be done only once per graph.
"""
global _finalized
tfutil.assert_tf_initialized()
if _finalized:
return None
_finalized = True
tfutil.init_uninitialized_vars([var for vars_list in _vars.values() for var in vars_list])
# Create summary ops.
with tf.device(None), tf.control_dependencies(None):
for name, vars_list in _vars.items():
name_id = name.replace("/", "_")
with tfutil.absolute_name_scope("Autosummary/" + name_id):
moments = tf.add_n(vars_list)
moments /= moments[0]
with tf.control_dependencies([moments]): # read before resetting
reset_ops = [tf.assign(var, tf.zeros(3, dtype=_dtype)) for var in vars_list]
with tf.name_scope(None), tf.control_dependencies(reset_ops): # reset before reporting
mean = moments[1]
std = tf.sqrt(moments[2] - tf.square(moments[1]))
tf.summary.scalar(name, mean)
if enable_custom_scalars:
tf.summary.scalar("xCustomScalars/" + name + "/margin_lo", mean - std)
tf.summary.scalar("xCustomScalars/" + name + "/margin_hi", mean + std)
# Setup layout for custom scalars.
layout = None
if enable_custom_scalars:
cat_dict = OrderedDict()
for series_name in sorted(_vars.keys()):
p = series_name.split("/")
cat = p[0] if len(p) >= 2 else ""
chart = "/".join(p[1:-1]) if len(p) >= 3 else p[-1]
if cat not in cat_dict:
cat_dict[cat] = OrderedDict()
if chart not in cat_dict[cat]:
cat_dict[cat][chart] = []
cat_dict[cat][chart].append(series_name)
categories = []
for cat_name, chart_dict in cat_dict.items():
charts = []
for chart_name, series_names in chart_dict.items():
series = []
for series_name in series_names:
series.append(layout_pb2.MarginChartContent.Series(
value=series_name,
lower="xCustomScalars/" + series_name + "/margin_lo",
upper="xCustomScalars/" + series_name + "/margin_hi"))
margin = layout_pb2.MarginChartContent(series=series)
charts.append(layout_pb2.Chart(title=chart_name, margin=margin))
categories.append(layout_pb2.Category(title=cat_name, chart=charts))
layout = summary_lib.custom_scalar_pb(layout_pb2.Layout(category=categories))
return layout
def save_summaries(file_writer, global_step=None):
"""Call FileWriter.add_summary() with all summaries in the default graph,
automatically finalizing and merging them on the first call.
"""
global _merge_op
tfutil.assert_tf_initialized()
if _merge_op is None:
layout = finalize_autosummaries()
if layout is not None:
file_writer.add_summary(layout)
with tf.device(None), tf.control_dependencies(None):
_merge_op = tf.summary.merge_all()
file_writer.add_summary(_merge_op.eval(), global_step)
================================================
FILE: stylegan_human/dnnlib/tflib/custom_ops.py
================================================
# Copyright (c) SenseTime Research. All rights reserved.
# Copyright (c) 2019, NVIDIA Corporation. All rights reserved.
#
# This work is made available under the Nvidia Source Code License-NC.
# To view a copy of this license, visit
# https://nvlabs.github.io/stylegan2/license.html
"""TensorFlow custom ops builder.
"""
import os
import re
import uuid
import hashlib
import tempfile
import shutil
import tensorflow as tf
from tensorflow.python.client import device_lib # pylint: disable=no-name-in-module
#----------------------------------------------------------------------------
# Global options.
cuda_cache_path = os.path.join(os.path.dirname(__file__), '_cudacache')
cuda_cache_version_tag = 'v1'
do_not_hash_included_headers = False # Speed up compilation by assuming that headers included by the CUDA code never change. Unsafe!
verbose = True # Print status messages to stdout.
compiler_bindir_search_path = [
'C:/Program Files (x86)/Microsoft Visual Studio/2017/Community/VC/Tools/MSVC/14.14.26428/bin/Hostx64/x64',
'C:/Program Files (x86)/Microsoft Visual Studio/2019/Community/VC/Tools/MSVC/14.23.28105/bin/Hostx64/x64',
'C:/Program Files (x86)/Microsoft Visual Studio 14.0/vc/bin',
]
#----------------------------------------------------------------------------
# Internal helper funcs.
def _find_compiler_bindir():
for compiler_path in compiler_bindir_search_path:
if os.path.isdir(compiler_path):
return compiler_path
return None
def _get_compute_cap(device):
caps_str = device.physical_device_desc
m = re.search('compute capability: (\\d+).(\\d+)', caps_str)
major = m.group(1)
minor = m.group(2)
return (major, minor)
def _get_cuda_gpu_arch_string():
gpus = [x for x in device_lib.list_local_devices() if x.device_type == 'GPU']
if len(gpus) == 0:
raise RuntimeError('No GPU devices found')
(major, minor) = _get_compute_cap(gpus[0])
return 'sm_%s%s' % (major, minor)
def _run_cmd(cmd):
with os.popen(cmd) as pipe:
output = pipe.read()
status = pipe.close()
if status is not None:
raise RuntimeError('NVCC returned an error. See below for full command line and output log:\n\n%s\n\n%s' % (cmd, output))
def _prepare_nvcc_cli(opts):
cmd = 'nvcc ' + opts.strip()
cmd += ' --disable-warnings'
cmd += ' --include-path "%s"' % tf.sysconfig.get_include()
cmd += ' --include-path "%s"' % os.path.join(tf.sysconfig.get_include(), 'external', 'protobuf_archive', 'src')
cmd += ' --include-path "%s"' % os.path.join(tf.sysconfig.get_include(), 'external', 'com_google_absl')
cmd += ' --include-path "%s"' % os.path.join(tf.sysconfig.get_include(), 'external', 'eigen_archive')
compiler_bindir = _find_compiler_bindir()
if compiler_bindir is None:
# Require that _find_compiler_bindir succeeds on Windows. Allow
# nvcc to use whatever is the default on Linux.
if os.name == 'nt':
raise RuntimeError('Could not find MSVC/GCC/CLANG installation on this computer. Check compiler_bindir_search_path list in "%s".' % __file__)
else:
cmd += ' --compiler-bindir "%s"' % compiler_bindir
cmd += ' 2>&1'
return cmd
#----------------------------------------------------------------------------
# Main entry point.
_plugin_cache = dict()
def get_plugin(cuda_file):
cuda_file_base = os.path.basename(cuda_file)
cuda_file_name, cuda_file_ext = os.path.splitext(cuda_file_base)
# Already in cache?
if cuda_file in _plugin_cache:
return _plugin_cache[cuda_file]
# Setup plugin.
if verbose:
print('Setting up TensorFlow plugin "%s": ' % cuda_file_base, end='', flush=True)
try:
# Hash CUDA source.
md5 = hashlib.md5()
with open(cuda_file, 'rb') as f:
md5.update(f.read())
md5.update(b'\n')
# Hash headers included by the CUDA code by running it through the preprocessor.
if not do_not_hash_included_headers:
if verbose:
print('Preprocessing... ', end='', flush=True)
with tempfile.TemporaryDirectory() as tmp_dir:
tmp_file = os.path.join(tmp_dir, cuda_file_name + '_tmp' + cuda_file_ext)
_run_cmd(_prepare_nvcc_cli('"%s" --preprocess -o "%s" --keep --keep-dir "%s"' % (cuda_file, tmp_file, tmp_dir)))
with open(tmp_file, 'rb') as f:
bad_file_str = ('"' + cuda_file.replace('\\', '/') + '"').encode('utf-8') # __FILE__ in error check macros
good_file_str = ('"' + cuda_file_base + '"').encode('utf-8')
for ln in f:
if not ln.startswith(b'# ') and not ln.startswith(b'#line '): # ignore line number pragmas
ln = ln.replace(bad_file_str, good_file_str)
md5.update(ln)
md5.update(b'\n')
# Select compiler options.
compile_opts = ''
if os.name == 'nt':
compile_opts += '"%s"' % os.path.join(tf.sysconfig.get_lib(), 'python', '_pywrap_tensorflow_internal.lib')
elif os.name == 'posix':
compile_opts += '"%s"' % os.path.join(tf.sysconfig.get_lib(), 'python', '_pywrap_tensorflow_internal.so')
compile_opts += ' --compiler-options \'-fPIC -D_GLIBCXX_USE_CXX11_ABI=0\''
else:
assert False # not Windows or Linux, w00t?
compile_opts += ' --gpu-architecture=%s' % _get_cuda_gpu_arch_string()
compile_opts += ' --use_fast_math'
nvcc_cmd = _prepare_nvcc_cli(compile_opts)
# Hash build configuration.
md5.update(('nvcc_cmd: ' + nvcc_cmd).encode('utf-8') + b'\n')
md5.update(('tf.VERSION: ' + tf.VERSION).encode('utf-8') + b'\n')
md5.update(('cuda_cache_version_tag: ' + cuda_cache_version_tag).encode('utf-8') + b'\n')
# Compile if not already compiled.
bin_file_ext = '.dll' if os.name == 'nt' else '.so'
bin_file = os.path.join(cuda_cache_path, cuda_file_name + '_' + md5.hexdigest() + bin_file_ext)
if not os.path.isfile(bin_file):
if verbose:
print('Compiling... ', end='', flush=True)
with tempfile.TemporaryDirectory() as tmp_dir:
tmp_file = os.path.join(tmp_dir, cuda_file_name + '_tmp' + bin_file_ext)
_run_cmd(nvcc_cmd + ' "%s" --shared -o "%s" --keep --keep-dir "%s"' % (cuda_file, tmp_file, tmp_dir))
os.makedirs(cuda_cache_path, exist_ok=True)
intermediate_file = os.path.join(cuda_cache_path, cuda_file_name + '_' + uuid.uuid4().hex + '_tmp' + bin_file_ext)
shutil.copyfile(tmp_file, intermediate_file)
os.rename(intermediate_file, bin_file) # atomic
# Load.
if verbose:
print('Loading... ', end='', flush=True)
plugin = tf.load_op_library(bin_file)
# Add to cache.
_plugin_cache[cuda_file] = plugin
if verbose:
print('Done.', flush=True)
return plugin
except:
if verbose:
print('Failed!', flush=True)
raise
#----------------------------------------------------------------------------
================================================
FILE: stylegan_human/dnnlib/tflib/network.py
================================================
# Copyright (c) SenseTime Research. All rights reserved.
# Copyright (c) 2019, NVIDIA Corporation. All rights reserved.
#
# This work is made available under the Nvidia Source Code License-NC.
# To view a copy of this license, visit
# https://nvlabs.github.io/stylegan2/license.html
"""Helper for managing networks."""
import types
import inspect
import re
import uuid
import sys
import numpy as np
import tensorflow as tf
from collections import OrderedDict
from typing import Any, List, Tuple, Union
from . import tfutil
from .. import util
from .tfutil import TfExpression, TfExpressionEx
_import_handlers = [] # Custom import handlers for dealing with legacy data in pickle import.
_import_module_src = dict() # Source code for temporary modules created during pickle import.
def import_handler(handler_func):
"""Function decorator for declaring custom import handlers."""
_import_handlers.append(handler_func)
return handler_func
class Network:
"""Generic network abstraction.
Acts as a convenience wrapper for a parameterized network construction
function, providing several utility methods and convenient access to
the inputs/outputs/weights.
Network objects can be safely pickled and unpickled for long-term
archival purposes. The pickling works reliably as long as the underlying
network construction function is defined in a standalone Python module
that has no side effects or application-specific imports.
Args:
name: Network name. Used to select TensorFlow name and variable scopes.
func_name: Fully qualified name of the underlying network construction function, or a top-level function object.
static_kwargs: Keyword arguments to be passed in to the network construction function.
Attributes:
name: User-specified name, defaults to build func name if None.
scope: Unique TensorFlow scope containing template graph and variables, derived from the user-specified name.
static_kwargs: Arguments passed to the user-supplied build func.
components: Container for sub-networks. Passed to the build func, and retained between calls.
num_inputs: Number of input tensors.
num_outputs: Number of output tensors.
input_shapes: Input tensor shapes (NC or NCHW), including minibatch dimension.
output_shapes: Output tensor shapes (NC or NCHW), including minibatch dimension.
input_shape: Short-hand for input_shapes[0].
output_shape: Short-hand for output_shapes[0].
input_templates: Input placeholders in the template graph.
output_templates: Output tensors in the template graph.
input_names: Name string for each input.
output_names: Name string for each output.
own_vars: Variables defined by this network (local_name => var), excluding sub-networks.
vars: All variables (local_name => var).
trainables: All trainable variables (local_name => var).
var_global_to_local: Mapping from variable global names to local names.
"""
def __init__(self, name: str = None, func_name: Any = None, **static_kwargs):
tfutil.assert_tf_initialized()
assert isinstance(name, str) or name is None
assert func_name is not None
assert isinstance(func_name, str) or util.is_top_level_function(func_name)
assert util.is_pickleable(static_kwargs)
self._init_fields()
self.name = name
self.static_kwargs = util.EasyDict(static_kwargs)
# Locate the user-specified network build function.
if util.is_top_level_function(func_name):
func_name = util.get_top_level_function_name(func_name)
module, self._build_func_name = util.get_module_from_obj_name(func_name)
self._build_func = util.get_obj_from_module(module, self._build_func_name)
assert callable(self._build_func)
# Dig up source code for the module containing the build function.
self._build_module_src = _import_module_src.get(module, None)
if self._build_module_src is None:
self._build_module_src = inspect.getsource(module)
# Init TensorFlow graph.
self._init_graph()
self.reset_own_vars()
def _init_fields(self) -> None:
self.name = None
self.scope = None
self.static_kwargs = util.EasyDict()
self.components = util.EasyDict()
self.num_inputs = 0
self.num_outputs = 0
self.input_shapes = [[]]
self.output_shapes = [[]]
self.input_shape = []
self.output_shape = []
self.input_templates = []
self.output_templates = []
self.input_names = []
self.output_names = []
self.own_vars = OrderedDict()
self.vars = OrderedDict()
self.trainables = OrderedDict()
self.var_global_to_local = OrderedDict()
self._build_func = None # User-supplied build function that constructs the network.
self._build_func_name = None # Name of the build function.
self._build_module_src = None # Full source code of the module containing the build function.
self._run_cache = dict() # Cached graph data for Network.run().
def _init_graph(self) -> None:
# Collect inputs.
self.input_names = []
for param in inspect.signature(self._build_func).parameters.values():
if param.kind == param.POSITIONAL_OR_KEYWORD and param.default is param.empty:
self.input_names.append(param.name)
self.num_inputs = len(self.input_names)
assert self.num_inputs >= 1
# Choose name and scope.
if self.name is None:
self.name = self._build_func_name
assert re.match("^[A-Za-z0-9_.\\-]*$", self.name)
with tf.name_scope(None):
self.scope = tf.get_default_graph().unique_name(self.name, mark_as_used=True)
# Finalize build func kwargs.
build_kwargs = dict(self.static_kwargs)
build_kwargs["is_template_graph"] = True
build_kwargs["components"] = self.components
# Build template graph.
with tfutil.absolute_variable_scope(self.scope, reuse=False), tfutil.absolute_name_scope(self.scope): # ignore surrounding scopes
assert tf.get_variable_scope().name == self.scope
assert tf.get_default_graph().get_name_scope() == self.scope
with tf.control_dependencies(None): # ignore surrounding control dependencies
self.input_templates = [tf.placeholder(tf.float32, name=name) for name in self.input_names]
out_expr = self._build_func(*self.input_templates, **build_kwargs)
# Collect outputs.
assert tfutil.is_tf_expression(out_expr) or isinstance(out_expr, tuple)
self.output_templates = [out_expr] if tfutil.is_tf_expression(out_expr) else list(out_expr)
self.num_outputs = len(self.output_templates)
assert self.num_outputs >= 1
assert all(tfutil.is_tf_expression(t) for t in self.output_templates)
# Perform sanity checks.
if any(t.shape.ndims is None for t in self.input_templates):
raise ValueError("Network input shapes not defined. Please call x.set_shape() for each input.")
if any(t.shape.ndims is None for t in self.output_templates):
raise ValueError("Network output shapes not defined. Please call x.set_shape() where applicable.")
if any(not isinstance(comp, Network) for comp in self.components.values()):
raise ValueError("Components of a Network must be Networks themselves.")
if len(self.components) != len(set(comp.name for comp in self.components.values())):
raise ValueError("Components of a Network must have unique names.")
# List inputs and outputs.
self.input_shapes = [t.shape.as_list() for t in self.input_templates]
self.output_shapes = [t.shape.as_list() for t in self.output_templates]
self.input_shape = self.input_shapes[0]
self.output_shape = self.output_shapes[0]
self.output_names = [t.name.split("/")[-1].split(":")[0] for t in self.output_templates]
# List variables.
self.own_vars = OrderedDict((var.name[len(self.scope) + 1:].split(":")[0], var) for var in tf.global_variables(self.scope + "/"))
self.vars = OrderedDict(self.own_vars)
self.vars.update((comp.name + "/" + name, var) for comp in self.components.values() for name, var in comp.vars.items())
self.trainables = OrderedDict((name, var) for name, var in self.vars.items() if var.trainable)
self.var_global_to_local = OrderedDict((var.name.split(":")[0], name) for name, var in self.vars.items())
def reset_own_vars(self) -> None:
"""Re-initialize all variables of this network, excluding sub-networks."""
tfutil.run([var.initializer for var in self.own_vars.values()])
def reset_vars(self) -> None:
"""Re-initialize all variables of this network, including sub-networks."""
tfutil.run([var.initializer for var in self.vars.values()])
def reset_trainables(self) -> None:
"""Re-initialize all trainable variables of this network, including sub-networks."""
tfutil.run([var.initializer for var in self.trainables.values()])
def get_output_for(self, *in_expr: TfExpression, return_as_list: bool = False, **dynamic_kwargs) -> Union[TfExpression, List[TfExpression]]:
"""Construct TensorFlow expression(s) for the output(s) of this network, given the input expression(s)."""
assert len(in_expr) == self.num_inputs
assert not all(expr is None for expr in in_expr)
# Finalize build func kwargs.
build_kwargs = dict(self.static_kwargs)
build_kwargs.update(dynamic_kwargs)
build_kwargs["is_template_graph"] = False
build_kwargs["components"] = self.components
# Build TensorFlow graph to evaluate the network.
with tfutil.absolute_variable_scope(self.scope, reuse=True), tf.name_scope(self.name):
assert tf.get_variable_scope().name == self.scope
valid_inputs = [expr for expr in in_expr if expr is not None]
final_inputs = []
for expr, name, shape in zip(in_expr, self.input_names, self.input_shapes):
if expr is not None:
expr = tf.identity(expr, name=name)
else:
expr = tf.zeros([tf.shape(valid_inputs[0])[0]] + shape[1:], name=name)
final_inputs.append(expr)
out_expr = self._build_func(*final_inputs, **build_kwargs)
# Propagate input shapes back to the user-specified expressions.
for expr, final in zip(in_expr, final_inputs):
if isinstance(expr, tf.Tensor):
expr.set_shape(final.shape)
# Express outputs in the desired format.
assert tfutil.is_tf_expression(out_expr) or isinstance(out_expr, tuple)
if return_as_list:
out_expr = [out_expr] if tfutil.is_tf_expression(out_expr) else list(out_expr)
return out_expr
def get_var_local_name(self, var_or_global_name: Union[TfExpression, str]) -> str:
"""Get the local name of a given variable, without any surrounding name scopes."""
assert tfutil.is_tf_expression(var_or_global_name) or isinstance(var_or_global_name, str)
global_name = var_or_global_name if isinstance(var_or_global_name, str) else var_or_global_name.name
return self.var_global_to_local[global_name]
def find_var(self, var_or_local_name: Union[TfExpression, str]) -> TfExpression:
"""Find variable by local or global name."""
assert tfutil.is_tf_expression(var_or_local_name) or isinstance(var_or_local_name, str)
return self.vars[var_or_local_name] if isinstance(var_or_local_name, str) else var_or_local_name
def get_var(self, var_or_local_name: Union[TfExpression, str]) -> np.ndarray:
"""Get the value of a given variable as NumPy array.
Note: This method is very inefficient -- prefer to use tflib.run(list_of_vars) whenever possible."""
return self.find_var(var_or_local_name).eval()
def set_var(self, var_or_local_name: Union[TfExpression, str], new_value: Union[int, float, np.ndarray]) -> None:
"""Set the value of a given variable based on the given NumPy array.
Note: This method is very inefficient -- prefer to use tflib.set_vars() whenever possible."""
tfutil.set_vars({self.find_var(var_or_local_name): new_value})
def __getstate__(self) -> dict:
"""Pickle export."""
state = dict()
state["version"] = 4
state["name"] = self.name
state["static_kwargs"] = dict(self.static_kwargs)
state["components"] = dict(self.components)
state["build_module_src"] = self._build_module_src
state["build_func_name"] = self._build_func_name
state["variables"] = list(zip(self.own_vars.keys(), tfutil.run(list(self.own_vars.values()))))
return state
def __setstate__(self, state: dict) -> None:
"""Pickle import."""
# pylint: disable=attribute-defined-outside-init
tfutil.assert_tf_initialized()
self._init_fields()
# Execute custom import handlers.
for handler in _import_handlers:
state = handler(state)
# Set basic fields.
assert state["version"] in [2, 3, 4]
self.name = state["name"]
self.static_kwargs = util.EasyDict(state["static_kwargs"])
self.components = util.EasyDict(state.get("components", {}))
self._build_module_src = state["build_module_src"]
self._build_func_name = state["build_func_name"]
# Create temporary module from the imported source code.
module_name = "_tflib_network_import_" + uuid.uuid4().hex
module = types.ModuleType(module_name)
sys.modules[module_name] = module
_import_module_src[module] = self._build_module_src
exec(self._build_module_src, module.__dict__) # pylint: disable=exec-used
# Locate network build function in the temporary module.
self._build_func = util.get_obj_from_module(module, self._build_func_name)
assert callable(self._build_func)
# Init TensorFlow graph.
self._init_graph()
self.reset_own_vars()
tfutil.set_vars({self.find_var(name): value for name, value in state["variables"]})
def clone(self, name: str = None, **new_static_kwargs) -> "Network":
"""Create a clone of this network with its own copy of the variables."""
# pylint: disable=protected-access
net = object.__new__(Network)
net._init_fields()
net.name = name if name is not None else self.name
net.static_kwargs = util.EasyDict(self.static_kwargs)
net.static_kwargs.update(new_static_kwargs)
net._build_module_src = self._build_module_src
net._build_func_name = self._build_func_name
net._build_func = self._build_func
net._init_graph()
net.copy_vars_from(self)
return net
def copy_own_vars_from(self, src_net: "Network") -> None:
"""Copy the values of all variables from the given network, excluding sub-networks."""
names = [name for name in self.own_vars.keys() if name in src_net.own_vars]
tfutil.set_vars(tfutil.run({self.vars[name]: src_net.vars[name] for name in names}))
def copy_vars_from(self, src_net: "Network") -> None:
"""Copy the values of all variables from the given network, including sub-networks."""
names = [name for name in self.vars.keys() if name in src_net.vars]
tfutil.set_vars(tfutil.run({self.vars[name]: src_net.vars[name] for name in names}))
def copy_trainables_from(self, src_net: "Network") -> None:
"""Copy the values of all trainable variables from the given network, including sub-networks."""
names = [name for name in self.trainables.keys() if name in src_net.trainables]
tfutil.set_vars(tfutil.run({self.vars[name]: src_net.vars[name] for name in names}))
def convert(self, new_func_name: str, new_name: str = None, **new_static_kwargs) -> "Network":
"""Create new network with the given parameters, and copy all variables from this network."""
if new_name is None:
new_name = self.name
static_kwargs = dict(self.static_kwargs)
static_kwargs.update(new_static_kwargs)
net = Network(name=new_name, func_name=new_func_name, **static_kwargs)
net.copy_vars_from(self)
return net
def setup_as_moving_average_of(self, src_net: "Network", beta: TfExpressionEx = 0.99, beta_nontrainable: TfExpressionEx = 0.0) -> tf.Operation:
"""Construct a TensorFlow op that updates the variables of this network
to be slightly closer to those of the given network."""
with tfutil.absolute_name_scope(self.scope + "/_MovingAvg"):
ops = []
for name, var in self.vars.items():
if name in src_net.vars:
cur_beta = beta if name in self.trainables else beta_nontrainable
new_value = tfutil.lerp(src_net.vars[name], var, cur_beta)
ops.append(var.assign(new_value))
return tf.group(*ops)
def run(self,
*in_arrays: Tuple[Union[np.ndarray, None], ...],
input_transform: dict = None,
output_transform: dict = None,
return_as_list: bool = False,
print_progress: bool = False,
minibatch_size: int = None,
num_gpus: int = 1,
assume_frozen: bool = False,
**dynamic_kwargs) -> Union[np.ndarray, Tuple[np.ndarray, ...], List[np.ndarray]]:
"""Run this network for the given NumPy array(s), and return the output(s) as NumPy array(s).
Args:
input_transform: A dict specifying a custom transformation to be applied to the input tensor(s) before evaluating the network.
The dict must contain a 'func' field that points to a top-level function. The function is called with the input
TensorFlow expression(s) as positional arguments. Any remaining fields of the dict will be passed in as kwargs.
output_transform: A dict specifying a custom transformation to be applied to the output tensor(s) after evaluating the network.
The dict must contain a 'func' field that points to a top-level function. The function is called with the output
TensorFlow expression(s) as positional arguments. Any remaining fields of the dict will be passed in as kwargs.
return_as_list: True = return a list of NumPy arrays, False = return a single NumPy array, or a tuple if there are multiple outputs.
print_progress: Print progress to the console? Useful for very large input arrays.
minibatch_size: Maximum minibatch size to use, None = disable batching.
num_gpus: Number of GPUs to use.
assume_frozen: Improve multi-GPU performance by assuming that the trainable parameters will remain changed between calls.
dynamic_kwargs: Additional keyword arguments to be passed into the network build function.
"""
assert len(in_arrays) == self.num_inputs
assert not all(arr is None for arr in in_arrays)
assert input_transform is None or util.is_top_level_function(input_transform["func"])
assert output_transform is None or util.is_top_level_function(output_transform["func"])
output_transform, dynamic_kwargs = _handle_legacy_output_transforms(output_transform, dynamic_kwargs)
num_items = in_arrays[0].shape[0]
if minibatch_size is None:
minibatch_size = num_items
# Construct unique hash key from all arguments that affect the TensorFlow graph.
key = dict(input_transform=input_transform, output_transform=output_transform, num_gpus=num_gpus, assume_frozen=assume_frozen, dynamic_kwargs=dynamic_kwargs)
def unwind_key(obj):
if isinstance(obj, dict):
return [(key, unwind_key(value)) for key, value in sorted(obj.items())]
if callable(obj):
return util.get_top_level_function_name(obj)
return obj
key = repr(unwind_key(key))
# Build graph.
if key not in self._run_cache:
with tfutil.absolute_name_scope(self.scope + "/_Run"), tf.control_dependencies(None):
with tf.device("/cpu:0"):
in_expr = [tf.placeholder(tf.float32, name=name) for name in self.input_names]
in_split = list(zip(*[tf.split(x, num_gpus) for x in in_expr]))
out_split = []
for gpu in range(num_gpus):
with tf.device("/gpu:%d" % gpu):
net_gpu = self.clone() if assume_frozen else self
in_gpu = in_split[gpu]
if input_transform is not None:
in_kwargs = dict(input_transform)
in_gpu = in_kwargs.pop("func")(*in_gpu, **in_kwargs)
in_gpu = [in_gpu] if tfutil.is_tf_expression(in_gpu) else list(in_gpu)
assert len(in_gpu) == self.num_inputs
out_gpu = net_gpu.get_output_for(*in_gpu, return_as_list=True, **dynamic_kwargs)
if output_transform is not None:
out_kwargs = dict(output_transform)
out_gpu = out_kwargs.pop("func")(*out_gpu, **out_kwargs)
out_gpu = [out_gpu] if tfutil.is_tf_expression(out_gpu) else list(out_gpu)
assert len(out_gpu) == self.num_outputs
out_split.append(out_gpu)
with tf.device("/cpu:0"):
out_expr = [tf.concat(outputs, axis=0) for outputs in zip(*out_split)]
self._run_cache[key] = in_expr, out_expr
# Run minibatches.
in_expr, out_expr = self._run_cache[key]
out_arrays = [np.empty([num_items] + expr.shape.as_list()[1:], expr.dtype.name) for expr in out_expr]
for mb_begin in range(0, num_items, minibatch_size):
if print_progress:
print("\r%d / %d" % (mb_begin, num_items), end="")
mb_end = min(mb_begin + minibatch_size, num_items)
mb_num = mb_end - mb_begin
mb_in = [src[mb_begin : mb_end] if src is not None else np.zeros([mb_num] + shape[1:]) for src, shape in zip(in_arrays, self.input_shapes)]
mb_out = tf.get_default_session().run(out_expr, dict(zip(in_expr, mb_in)))
for dst, src in zip(out_arrays, mb_out):
dst[mb_begin: mb_end] = src
# Done.
if print_progress:
print("\r%d / %d" % (num_items, num_items))
if not return_as_list:
out_arrays = out_arrays[0] if len(out_arrays) == 1 else tuple(out_arrays)
return out_arrays
def list_ops(self) -> List[TfExpression]:
include_prefix = self.scope + "/"
exclude_prefix = include_prefix + "_"
ops = tf.get_default_graph().get_operations()
ops = [op for op in ops if op.name.startswith(include_prefix)]
ops = [op for op in ops if not op.name.startswith(exclude_prefix)]
return ops
def list_layers(self) -> List[Tuple[str, TfExpression, List[TfExpression]]]:
"""Returns a list of (layer_name, output_expr, trainable_vars) tuples corresponding to
individual layers of the network. Mainly intended to be used for reporting."""
layers = []
def recurse(scope, parent_ops, parent_vars, level):
# Ignore specific patterns.
if any(p in scope for p in ["/Shape", "/strided_slice", "/Cast", "/concat", "/Assign"]):
return
# Filter ops and vars by scope.
global_prefix = scope + "/"
local_prefix = global_prefix[len(self.scope) + 1:]
cur_ops = [op for op in parent_ops if op.name.startswith(global_prefix) or op.name == global_prefix[:-1]]
cur_vars = [(name, var) for name, var in parent_vars if name.startswith(local_prefix) or name == local_prefix[:-1]]
if not cur_ops and not cur_vars:
return
# Filter out all ops related to variables.
for var in [op for op in cur_ops if op.type.startswith("Variable")]:
var_prefix = var.name + "/"
cur_ops = [op for op in cur_ops if not op.name.startswith(var_prefix)]
# Scope does not contain ops as immediate children => recurse deeper.
contains_direct_ops = any("/" not in op.name[len(global_prefix):] and op.type not in ["Identity", "Cast", "Transpose"] for op in cur_ops)
if (level == 0 or not contains_direct_ops) and (len(cur_ops) + len(cur_vars)) > 1:
visited = set()
for rel_name in [op.name[len(global_prefix):] for op in cur_ops] + [name[len(local_prefix):] for name, _var in cur_vars]:
token = rel_name.split("/")[0]
if token not in visited:
recurse(global_prefix + token, cur_ops, cur_vars, level + 1)
visited.add(token)
return
# Report layer.
layer_name = scope[len(self.scope) + 1:]
layer_output = cur_ops[-1].outputs[0] if cur_ops else cur_vars[-1][1]
layer_trainables = [var for _name, var in cur_vars if var.trainable]
layers.append((layer_name, layer_output, layer_trainables))
recurse(self.scope, self.list_ops(), list(self.vars.items()), 0)
return layers
def print_layers(self, title: str = None, hide_layers_with_no_params: bool = False) -> None:
"""Print a summary table of the network structure."""
rows = [[title if title is not None else self.name, "Params", "OutputShape", "WeightShape"]]
rows += [["---"] * 4]
total_params = 0
for layer_name, layer_output, layer_trainables in self.list_layers():
num_params = sum(int(np.prod(var.shape.as_list())) for var in layer_trainables)
weights = [var for var in layer_trainables if var.name.endswith("/weight:0")]
weights.sort(key=lambda x: len(x.name))
if len(weights) == 0 and len(layer_trainables) == 1:
weights = layer_trainables
total_params += num_params
if not hide_layers_with_no_params or num_params != 0:
num_params_str = str(num_params) if num_params > 0 else "-"
output_shape_str = str(layer_output.shape)
weight_shape_str = str(weights[0].shape) if len(weights) >= 1 else "-"
rows += [[layer_name, num_params_str, output_shape_str, weight_shape_str]]
rows += [["---"] * 4]
rows += [["Total", str(total_params), "", ""]]
widths = [max(len(cell) for cell in column) for column in zip(*rows)]
print()
for row in rows:
print(" ".join(cell + " " * (width - len(cell)) for cell, width in zip(row, widths)))
print()
def setup_weight_histograms(self, title: str = None) -> None:
"""Construct summary ops to include histograms of all trainable parameters in TensorBoard."""
if title is None:
title = self.name
with tf.name_scope(None), tf.device(None), tf.control_dependencies(None):
for local_name, var in self.trainables.items():
if "/" in local_name:
p = local_name.split("/")
name = title + "_" + p[-1] + "/" + "_".join(p[:-1])
else:
name = title + "_toplevel/" + local_name
tf.summary.histogram(name, var)
#----------------------------------------------------------------------------
# Backwards-compatible emulation of legacy output transformation in Network.run().
_print_legacy_warning = True
def _handle_legacy_output_transforms(output_transform, dynamic_kwargs):
global _print_legacy_warning
legacy_kwargs = ["out_mul", "out_add", "out_shrink", "out_dtype"]
if not any(kwarg in dynamic_kwargs for kwarg in legacy_kwargs):
return output_transform, dynamic_kwargs
if _print_legacy_warning:
_print_legacy_warning = False
print()
print("WARNING: Old-style output transformations in Network.run() are deprecated.")
print("Consider using 'output_transform=dict(func=tflib.convert_images_to_uint8)'")
print("instead of 'out_mul=127.5, out_add=127.5, out_dtype=np.uint8'.")
print()
assert output_transform is None
new_kwargs = dict(dynamic_kwargs)
new_transform = {kwarg: new_kwargs.pop(kwarg) for kwarg in legacy_kwargs if kwarg in dynamic_kwargs}
new_transform["func"] = _legacy_output_transform_func
return new_transform, new_kwargs
def _legacy_output_transform_func(*expr, out_mul=1.0, out_add=0.0, out_shrink=1, out_dtype=None):
if out_mul != 1.0:
expr = [x * out_mul for x in expr]
if out_add != 0.0:
expr = [x + out_add for x in expr]
if out_shrink > 1:
ksize = [1, 1, out_shrink, out_shrink]
expr = [tf.nn.avg_pool(x, ksize=ksize, strides=ksize, padding="VALID", data_format="NCHW") for x in expr]
if out_dtype is not None:
if tf.as_dtype(out_dtype).is_integer:
expr = [tf.round(x) for x in expr]
expr = [tf.saturate_cast(x, out_dtype) for x in expr]
return expr
================================================
FILE: stylegan_human/dnnlib/tflib/ops/__init__.py
================================================
# Copyright (c) SenseTime Research. All rights reserved.
# Copyright (c) 2019, NVIDIA Corporation. All rights reserved.
#
# This work is made available under the Nvidia Source Code License-NC.
# To view a copy of this license, visit
# https://nvlabs.github.io/stylegan2/license.html
# empty
================================================
FILE: stylegan_human/dnnlib/tflib/ops/fused_bias_act.cu
================================================
// Copyright (c) SenseTime Research. All rights reserved.
// Copyright (c) 2019, NVIDIA Corporation. All rights reserved.
//
// This work is made available under the Nvidia Source Code License-NC.
// To view a copy of this license, visit
// https://nvlabs.github.io/stylegan2/license.html
#define EIGEN_USE_GPU
#define __CUDA_INCLUDE_COMPILER_INTERNAL_HEADERS__
#include "tensorflow/core/framework/op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/shape_inference.h"
#include <stdio.h>
using namespace tensorflow;
using namespace tensorflow::shape_inference;
#define OP_CHECK_CUDA_ERROR(CTX, CUDA_CALL) do { cudaError_t err = CUDA_CALL; OP_REQUIRES(CTX, err == cudaSuccess, errors::Internal(cudaGetErrorName(err))); } while (false)
//------------------------------------------------------------------------
// CUDA kernel.
template <class T>
struct FusedBiasActKernelParams
{
const T* x; // [sizeX]
const T* b; // [sizeB] or NULL
const T* ref; // [sizeX] or NULL
T* y; // [sizeX]
int grad;
int axis;
int act;
float alpha;
float gain;
int sizeX;
int sizeB;
int stepB;
int loopX;
};
template <class T>
static __global__ void FusedBiasActKernel(const FusedBiasActKernelParams<T> p)
{
const float expRange = 80.0f;
const float halfExpRange = 40.0f;
const float seluScale = 1.0507009873554804934193349852946f;
const float seluAlpha = 1.6732632423543772848170429916717f;
// Loop over elements.
int xi = blockIdx.x * p.loopX * blockDim.x + threadIdx.x;
for (int loopIdx = 0; loopIdx < p.loopX && xi < p.sizeX; loopIdx++, xi += blockDim.x)
{
// Load and apply bias.
float x = (float)p.x[xi];
if (p.b)
x += (float)p.b[(xi / p.stepB) % p.sizeB];
float ref = (p.ref) ? (float)p.ref[xi] : 0.0f;
if (p.gain != 0.0f & p.act != 9)
ref /= p.gain;
// Evaluate activation func.
float y;
switch (p.act * 10 + p.grad)
{
// linear
default:
case 10: y = x; break;
case 11: y = x; break;
case 12: y = 0.0f; break;
// relu
case 20: y = (x > 0.0f) ? x : 0.0f; break;
case 21: y = (ref > 0.0f) ? x : 0.0f; break;
case 22: y = 0.0f; break;
// lrelu
case 30: y = (x > 0.0f) ? x : x * p.alpha; break;
case 31: y = (ref > 0.0f) ? x : x * p.alpha; break;
case 32: y = 0.0f; break;
// tanh
case 40: { float c = expf(x); float d = 1.0f / c; y = (x < -expRange) ? -1.0f : (x > expRange) ? 1.0f : (c - d) / (c + d); } break;
case 41: y = x * (1.0f - ref * ref); break;
case 42: y = x * (1.0f - ref * ref) * (-2.0f * ref); break;
// sigmoid
case 50: y = (x < -expRange) ? 0.0f : 1.0f / (expf(-x) + 1.0f); break;
case 51: y = x * ref * (1.0f - ref); break;
case 52: y = x * ref * (1.0f - ref) * (1.0f - 2.0f * ref); break;
// elu
case 60: y = (x >= 0.0f) ? x : expf(x) - 1.0f; break;
case 61: y = (ref >= 0.0f) ? x : x * (ref + 1.0f); break;
case 62: y = (ref >= 0.0f) ? 0.0f : x * (ref + 1.0f); break;
// selu
case 70: y = (x >= 0.0f) ? seluScale * x : (seluScale * seluAlpha) * (expf(x) - 1.0f); break;
case 71: y = (ref >= 0.0f) ? x * seluScale : x * (ref + seluScale * seluAlpha); break;
case 72: y = (ref >= 0.0f) ? 0.0f : x * (ref + seluScale * seluAlpha); break;
// softplus
case 80: y = (x > expRange) ? x : logf(expf(x) + 1.0f); break;
case 81: y = x * (1.0f - expf(-ref)); break;
case 82: { float c = expf(-ref); y = x * c * (1.0f - c); } break;
// swish
case 90: y = (x < -expRange) ? 0.0f : x / (expf(-x) + 1.0f); break;
case 91: { float c = expf(ref); float d = c + 1.0f; y = (ref > halfExpRange) ? x : x * c * (ref + d) / (d * d); } break;
case 92: { float c = expf(ref); float d = c + 1.0f; y = (ref > halfExpRange) ? 0.0f : x * c * (ref * (2.0f - d) + 2.0f * d) / (d * d * d); } break;
}
// Apply gain and store.
p.y[xi] = (T)(y * p.gain);
}
}
//------------------------------------------------------------------------
// TensorFlow op.
template <class T>
struct FusedBiasActOp : public OpKernel
{
FusedBiasActKernelParams<T> m_attribs;
FusedBiasActOp(OpKernelConstruction* ctx) : OpKernel(ctx)
{
memset(&m_attribs, 0, sizeof(m_attribs));
OP_REQUIRES_OK(ctx, ctx->GetAttr("grad", &m_attribs.grad));
OP_REQUIRES_OK(ctx, ctx->GetAttr("axis", &m_attribs.axis));
OP_REQUIRES_OK(ctx, ctx->GetAttr("act", &m_attribs.act));
OP_REQUIRES_OK(ctx, ctx->GetAttr("alpha", &m_attribs.alpha));
OP_REQUIRES_OK(ctx, ctx->GetAttr("gain", &m_attribs.gain));
OP_REQUIRES(ctx, m_attribs.grad >= 0, errors::InvalidArgument("grad must be non-negative"));
OP_REQUIRES(ctx, m_attribs.axis >= 0, errors::InvalidArgument("axis must be non-negative"));
OP_REQUIRES(ctx, m_attribs.act >= 0, errors::InvalidArgument("act must be non-negative"));
}
void Compute(OpKernelContext* ctx)
{
FusedBiasActKernelParams<T> p = m_attribs;
cudaStream_t stream = ctx->eigen_device<Eigen::GpuDevice>().stream();
const Tensor& x = ctx->input(0); // [...]
const Tensor& b = ctx->input(1); // [sizeB] or [0]
const Tensor& ref = ctx->input(2); // x.shape or [0]
p.x = x.flat<T>().data();
p.b = (b.NumElements()) ? b.flat<T>().data() : NULL;
p.ref = (ref.NumElements()) ? ref.flat<T>().data() : NULL;
OP_REQUIRES(ctx, b.NumElements() == 0 || m_attribs.axis < x.dims(), errors::InvalidArgument("axis out of bounds"));
OP_REQUIRES(ctx, b.dims() == 1, errors::InvalidArgument("b must have rank 1"));
OP_REQUIRES(ctx, b.NumElements() == 0 || b.NumElements() == x.dim_size(m_attribs.axis), errors::InvalidArgument("b has wrong number of elements"));
OP_REQUIRES(ctx, ref.NumElements() == ((p.grad == 0) ? 0 : x.NumElements()), errors::InvalidArgument("ref has wrong number of elements"));
OP_REQUIRES(ctx, x.NumElements() <= kint32max, errors::InvalidArgument("x is too large"));
p.sizeX = (int)x.NumElements();
p.sizeB = (int)b.NumElements();
p.stepB = 1;
for (int i = m_attribs.axis + 1; i < x.dims(); i++)
p.stepB *= (int)x.dim_size(i);
Tensor* y = NULL; // x.shape
OP_REQUIRES_OK(ctx, ctx->allocate_output(0, x.shape(), &y));
p.y = y->flat<T>().data();
p.loopX = 4;
int blockSize = 4 * 32;
int gridSize = (p.sizeX - 1) / (p.loopX * blockSize) + 1;
void* args[] = {&p};
OP_CHECK_CUDA_ERROR(ctx, cudaLaunchKernel((void*)FusedBiasActKernel<T>, gridSize, blockSize, args, 0, stream));
}
};
REGISTER_OP("FusedBiasAct")
.Input ("x: T")
.Input ("b: T")
.Input ("ref: T")
.Output ("y: T")
.Attr ("T: {float, half}")
.Attr ("grad: int = 0")
.Attr ("axis: int = 1")
.Attr ("act: int = 0")
.Attr ("alpha: float = 0.0")
.Attr ("gain: float = 1.0");
REGISTER_KERNEL_BUILDER(Name("FusedBiasAct").Device(DEVICE_GPU).TypeConstraint<float>("T"), FusedBiasActOp<float>);
REGISTER_KERNEL_BUILDER(Name("FusedBiasAct").Device(DEVICE_GPU).TypeConstraint<Eigen::half>("T"), FusedBiasActOp<Eigen::half>);
//------------------------------------------------------------------------
================================================
FILE: stylegan_human/dnnlib/tflib/ops/fused_bias_act.py
================================================
# Copyright (c) SenseTime Research. All rights reserved.
# Copyright (c) 2019, NVIDIA Corporation. All rights reserved.
#
# This work is made available under the Nvidia Source Code License-NC.
# To view a copy of this license, visit
# https://nvlabs.github.io/stylegan2/license.html
"""Custom TensorFlow ops for efficient bias and activation."""
import os
import numpy as np
import tensorflow as tf
from .. import custom_ops
from ...util import EasyDict
def _get_plugin():
return custom_ops.get_plugin(os.path.splitext(__file__)[0] + '.cu')
#----------------------------------------------------------------------------
activation_funcs = {
'linear': EasyDict(func=lambda x, **_: x, def_alpha=None, def_gain=1.0, cuda_idx=1, ref='y', zero_2nd_grad=True),
'relu': EasyDict(func=lambda x, **_: tf.nn.relu(x), def_alpha=None, def_gain=np.sqrt(2), cuda_idx=2, ref='y', zero_2nd_grad=True),
'lrelu': EasyDict(func=lambda x, alpha, **_: tf.nn.leaky_relu(x, alpha), def_alpha=0.2, def_gain=np.sqrt(2), cuda_idx=3, ref='y', zero_2nd_grad=True),
'tanh': EasyDict(func=lambda x, **_: tf.nn.tanh(x), def_alpha=None, def_gain=1.0, cuda_idx=4, ref='y', zero_2nd_grad=False),
'sigmoid': EasyDict(func=lambda x, **_: tf.nn.sigmoid(x), def_alpha=None, def_gain=1.0, cuda_idx=5, ref='y', zero_2nd_grad=False),
'elu': EasyDict(func=lambda x, **_: tf.nn.elu(x), def_alpha=None, def_gain=1.0, cuda_idx=6, ref='y', zero_2nd_grad=False),
'selu': EasyDict(func=lambda x, **_: tf.nn.selu(x), def_alpha=None, def_gain=1.0, cuda_idx=7, ref='y', zero_2nd_grad=False),
'softplus': EasyDict(func=lambda x, **_: tf.nn.softplus(x), def_alpha=None, def_gain=1.0, cuda_idx=8, ref='y', zero_2nd_grad=False),
'swish': EasyDict(func=lambda x, **_: tf.nn.sigmoid(x) * x, def_alpha=None, def_gain=np.sqrt(2), cuda_idx=9, ref='x', zero_2nd_grad=False),
}
#------------------------------------------
gitextract_5nncgy8o/
├── .gitignore
├── Dockerfile
├── LICENSE.txt
├── README.md
├── dnnlib/
│ ├── __init__.py
│ └── util.py
├── environment.yml
├── gen_images.py
├── gradio_utils/
│ ├── __init__.py
│ └── utils.py
├── gui_utils/
│ ├── __init__.py
│ ├── gl_utils.py
│ ├── glfw_window.py
│ ├── imgui_utils.py
│ ├── imgui_window.py
│ └── text_utils.py
├── legacy.py
├── requirements.txt
├── stylegan_human/
│ ├── .gitignore
│ ├── PP_HumanSeg/
│ │ ├── deploy/
│ │ │ └── infer.py
│ │ ├── export_model/
│ │ │ └── download_export_model.py
│ │ └── pretrained_model/
│ │ └── download_pretrained_model.py
│ ├── README.md
│ ├── __init__.py
│ ├── alignment.py
│ ├── bg_white.py
│ ├── dnnlib/
│ │ ├── __init__.py
│ │ ├── tflib/
│ │ │ ├── __init__.py
│ │ │ ├── autosummary.py
│ │ │ ├── custom_ops.py
│ │ │ ├── network.py
│ │ │ ├── ops/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── fused_bias_act.cu
│ │ │ │ ├── fused_bias_act.py
│ │ │ │ ├── upfirdn_2d.cu
│ │ │ │ └── upfirdn_2d.py
│ │ │ ├── optimizer.py
│ │ │ └── tfutil.py
│ │ └── util.py
│ ├── docs/
│ │ └── Dataset.md
│ ├── edit/
│ │ ├── __init__.py
│ │ ├── edit_config.py
│ │ └── edit_helper.py
│ ├── edit.py
│ ├── environment.yml
│ ├── generate.py
│ ├── insetgan.py
│ ├── interpolation.py
│ ├── latent_direction/
│ │ └── ss_statics/
│ │ ├── bottom_length_statis/
│ │ │ ├── 3/
│ │ │ │ └── statis.csv
│ │ │ ├── 4/
│ │ │ │ └── statis.csv
│ │ │ └── 5/
│ │ │ └── statis.csv
│ │ └── upper_length_statis/
│ │ └── 5/
│ │ └── statis.csv
│ ├── legacy.py
│ ├── openpose/
│ │ ├── model/
│ │ │ └── .gitkeep
│ │ └── src/
│ │ ├── __init__.py
│ │ ├── body.py
│ │ ├── model.py
│ │ └── util.py
│ ├── pti/
│ │ ├── pti_configs/
│ │ │ ├── __init__.py
│ │ │ ├── global_config.py
│ │ │ ├── hyperparameters.py
│ │ │ └── paths_config.py
│ │ ├── pti_models/
│ │ │ ├── __init__.py
│ │ │ └── e4e/
│ │ │ ├── __init__.py
│ │ │ ├── encoders/
│ │ │ │ ├── __init__.py
│ │ │ │ ├── helpers.py
│ │ │ │ ├── model_irse.py
│ │ │ │ └── psp_encoders.py
│ │ │ ├── latent_codes_pool.py
│ │ │ ├── psp.py
│ │ │ └── stylegan2/
│ │ │ ├── __init__.py
│ │ │ ├── model.py
│ │ │ └── op/
│ │ │ ├── __init__.py
│ │ │ ├── fused_act.py
│ │ │ ├── fused_bias_act.cpp
│ │ │ ├── fused_bias_act_kernel.cu
│ │ │ ├── upfirdn2d.cpp
│ │ │ ├── upfirdn2d.py
│ │ │ └── upfirdn2d_kernel.cu
│ │ └── training/
│ │ ├── __init__.py
│ │ ├── coaches/
│ │ │ ├── __init__.py
│ │ │ ├── base_coach.py
│ │ │ ├── localitly_regulizer.py
│ │ │ ├── multi_id_coach.py
│ │ │ └── single_id_coach.py
│ │ └── projectors/
│ │ ├── __init__.py
│ │ ├── w_plus_projector.py
│ │ └── w_projector.py
│ ├── run_pti.py
│ ├── style_mixing.py
│ ├── stylemixing_video.py
│ ├── torch_utils/
│ │ ├── __init__.py
│ │ ├── custom_ops.py
│ │ ├── misc.py
│ │ ├── models.py
│ │ ├── models_face.py
│ │ ├── op_edit/
│ │ │ ├── __init__.py
│ │ │ ├── fused_act.py
│ │ │ ├── fused_bias_act.cpp
│ │ │ ├── fused_bias_act_kernel.cu
│ │ │ ├── upfirdn2d.cpp
│ │ │ ├── upfirdn2d.py
│ │ │ └── upfirdn2d_kernel.cu
│ │ ├── ops/
│ │ │ ├── __init__.py
│ │ │ ├── bias_act.cpp
│ │ │ ├── bias_act.cu
│ │ │ ├── bias_act.h
│ │ │ ├── bias_act.py
│ │ │ ├── conv2d_gradfix.py
│ │ │ ├── conv2d_resample.py
│ │ │ ├── filtered_lrelu.cpp
│ │ │ ├── filtered_lrelu.cu
│ │ │ ├── filtered_lrelu.h
│ │ │ ├── filtered_lrelu.py
│ │ │ ├── filtered_lrelu_ns.cu
│ │ │ ├── filtered_lrelu_rd.cu
│ │ │ ├── filtered_lrelu_wr.cu
│ │ │ ├── fma.py
│ │ │ ├── grid_sample_gradfix.py
│ │ │ ├── upfirdn2d.cpp
│ │ │ ├── upfirdn2d.cu
│ │ │ ├── upfirdn2d.h
│ │ │ └── upfirdn2d.py
│ │ ├── persistence.py
│ │ └── training_stats.py
│ ├── training/
│ │ ├── __init__.py
│ │ ├── augment.py
│ │ ├── dataset.py
│ │ ├── loss.py
│ │ ├── networks_stylegan2.py
│ │ ├── networks_stylegan3.py
│ │ └── training_loop.py
│ ├── training_scripts/
│ │ ├── sg2/
│ │ │ ├── train.py
│ │ │ └── training/
│ │ │ ├── dataset.py
│ │ │ └── networks.py
│ │ └── sg3/
│ │ ├── train.py
│ │ └── training/
│ │ ├── dataset.py
│ │ ├── networks_stylegan2.py
│ │ └── networks_stylegan3.py
│ └── utils/
│ ├── ImagesDataset.py
│ ├── __init__.py
│ ├── data_utils.py
│ ├── face_alignment.py
│ ├── log_utils.py
│ ├── models_utils.py
│ └── util.py
├── torch_utils/
│ ├── __init__.py
│ ├── custom_ops.py
│ ├── misc.py
│ ├── ops/
│ │ ├── __init__.py
│ │ ├── bias_act.cpp
│ │ ├── bias_act.cu
│ │ ├── bias_act.h
│ │ ├── bias_act.py
│ │ ├── conv2d_gradfix.py
│ │ ├── conv2d_resample.py
│ │ ├── filtered_lrelu.cpp
│ │ ├── filtered_lrelu.cu
│ │ ├── filtered_lrelu.h
│ │ ├── filtered_lrelu.py
│ │ ├── filtered_lrelu_ns.cu
│ │ ├── filtered_lrelu_rd.cu
│ │ ├── filtered_lrelu_wr.cu
│ │ ├── fma.py
│ │ ├── grid_sample_gradfix.py
│ │ ├── upfirdn2d.cpp
│ │ ├── upfirdn2d.cu
│ │ ├── upfirdn2d.h
│ │ └── upfirdn2d.py
│ ├── persistence.py
│ └── training_stats.py
├── training/
│ ├── __init__.py
│ ├── augment.py
│ ├── dataset.py
│ ├── loss.py
│ ├── networks_stylegan2.py
│ ├── networks_stylegan3.py
│ └── training_loop.py
├── visualizer_drag.py
├── visualizer_drag_gradio.py
└── viz/
├── __init__.py
├── capture_widget.py
├── drag_widget.py
├── latent_widget.py
├── pickle_widget.py
└── renderer.py
SYMBOL INDEX (1360 symbols across 120 files)
FILE: dnnlib/util.py
class EasyDict (line 40) | class EasyDict(dict):
method __getattr__ (line 43) | def __getattr__(self, name: str) -> Any:
method __setattr__ (line 49) | def __setattr__(self, name: str, value: Any) -> None:
method __delattr__ (line 52) | def __delattr__(self, name: str) -> None:
class Logger (line 56) | class Logger(object):
method __init__ (line 59) | def __init__(self, file_name: str = None, file_mode: str = "w", should...
method __enter__ (line 72) | def __enter__(self) -> "Logger":
method __exit__ (line 75) | def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> N...
method write (line 78) | def write(self, text: Union[str, bytes]) -> None:
method flush (line 93) | def flush(self) -> None:
method close (line 100) | def close(self) -> None:
function set_cache_dir (line 120) | def set_cache_dir(path: str) -> None:
function make_cache_dir_path (line 124) | def make_cache_dir_path(*paths: str) -> str:
function format_time (line 139) | def format_time(seconds: Union[int, float]) -> str:
function format_time_brief (line 153) | def format_time_brief(seconds: Union[int, float]) -> str:
function ask_yes_no (line 167) | def ask_yes_no(question: str) -> bool:
function tuple_product (line 177) | def tuple_product(t: Tuple) -> Any:
function get_dtype_and_ctype (line 201) | def get_dtype_and_ctype(type_obj: Any) -> Tuple[np.dtype, Any]:
function is_pickleable (line 224) | def is_pickleable(obj: Any) -> bool:
function get_module_from_obj_name (line 236) | def get_module_from_obj_name(obj_name: str) -> Tuple[types.ModuleType, s...
function get_obj_from_module (line 277) | def get_obj_from_module(module: types.ModuleType, obj_name: str) -> Any:
function get_obj_by_name (line 287) | def get_obj_by_name(name: str) -> Any:
function call_func_by_name (line 293) | def call_func_by_name(*args, func_name: str = None, **kwargs) -> Any:
function construct_class_by_name (line 301) | def construct_class_by_name(*args, class_name: str = None, **kwargs) -> ...
function get_module_dir_by_obj_name (line 306) | def get_module_dir_by_obj_name(obj_name: str) -> str:
function is_top_level_function (line 312) | def is_top_level_function(obj: Any) -> bool:
function get_top_level_function_name (line 317) | def get_top_level_function_name(obj: Any) -> str:
function list_dir_recursively_with_ignore (line 329) | def list_dir_recursively_with_ignore(dir_path: str, ignores: List[str] =...
function copy_files_and_create_dirs (line 362) | def copy_files_and_create_dirs(files: List[Tuple[str, str]]) -> None:
function is_url (line 378) | def is_url(obj: Any, allow_file_urls: bool = False) -> bool:
function open_url (line 396) | def open_url(url: str, cache_dir: str = None, num_attempts: int = 10, ve...
FILE: gen_images.py
function parse_range (line 25) | def parse_range(s: Union[str, List]) -> List[int]:
function parse_vec2 (line 43) | def parse_vec2(s: Union[str, Tuple[float, float]]) -> Tuple[float, float]:
function make_transform (line 57) | def make_transform(translate: Tuple[float,float], angle: float):
function generate_images (line 80) | def generate_images(
FILE: gradio_utils/utils.py
class ImageMask (line 6) | class ImageMask(gr.components.Image):
method __init__ (line 13) | def __init__(self, **kwargs):
method preprocess (line 19) | def preprocess(self, x):
function get_valid_mask (line 33) | def get_valid_mask(mask: np.ndarray):
function draw_points_on_image (line 44) | def draw_points_on_image(image,
function draw_mask_on_image (line 112) | def draw_mask_on_image(image, mask):
function on_change_single_global_state (line 128) | def on_change_single_global_state(keys,
function get_latest_points_pair (line 149) | def get_latest_points_pair(points_dict):
FILE: gui_utils/gl_utils.py
function init_egl (line 20) | def init_egl():
function get_texture_format (line 77) | def get_texture_format(dtype, channels):
function prepare_texture_data (line 82) | def prepare_texture_data(image):
function draw_pixels (line 92) | def draw_pixels(image, *, pos=0, zoom=1, align=0, rint=True):
function read_pixels (line 115) | def read_pixels(width, height, *, pos=0, dtype='uint8', channels=3):
class Texture (line 129) | class Texture:
method __init__ (line 130) | def __init__(self, *, image=None, width=None, height=None, channels=No...
method delete (line 162) | def delete(self):
method __del__ (line 167) | def __del__(self):
method bind (line 174) | def bind(self):
method update (line 180) | def update(self, image):
method draw (line 193) | def draw(self, *, pos=0, zoom=1, align=0, rint=False, color=1, alpha=1...
method is_compatible (line 202) | def is_compatible(self, *, image=None, width=None, height=None, channe...
class Framebuffer (line 221) | class Framebuffer:
method __init__ (line 222) | def __init__(self, *, texture=None, width=None, height=None, channels=...
method delete (line 273) | def delete(self):
method __del__ (line 284) | def __del__(self):
method bind (line 291) | def bind(self):
method blit (line 301) | def blit(self, dst=None):
function draw_shape (line 309) | def draw_shape(vertices, *, mode=gl.GL_TRIANGLE_FAN, pos=0, size=1, colo...
function draw_arrow (line 336) | def draw_arrow(x1, y1, x2, y2, l=10, width=1.0):
function draw_rect (line 377) | def draw_rect(*, pos=0, pos2=None, size=None, align=0, rint=False, color...
function _setup_rect (line 394) | def _setup_rect(rx, ry):
function draw_circle (line 404) | def draw_circle(*, center=0, radius=100, hole=0, color=1, alpha=1):
function _setup_circle (line 410) | def _setup_circle(hole):
FILE: gui_utils/glfw_window.py
class GlfwWindow (line 16) | class GlfwWindow: # pylint: disable=too-many-public-methods
method __init__ (line 17) | def __init__(self, *, title='GlfwWindow', window_width=1920, window_he...
method close (line 45) | def close(self):
method __del__ (line 53) | def __del__(self):
method window_width (line 60) | def window_width(self):
method window_height (line 64) | def window_height(self):
method content_width (line 68) | def content_width(self):
method content_height (line 73) | def content_height(self):
method title_bar_height (line 78) | def title_bar_height(self):
method monitor_width (line 83) | def monitor_width(self):
method monitor_height (line 88) | def monitor_height(self):
method frame_delta (line 93) | def frame_delta(self):
method set_title (line 96) | def set_title(self, title):
method set_window_size (line 99) | def set_window_size(self, width, height):
method set_content_size (line 106) | def set_content_size(self, width, height):
method maximize (line 109) | def maximize(self):
method set_position (line 112) | def set_position(self, x, y):
method center (line 115) | def center(self):
method set_vsync (line 118) | def set_vsync(self, vsync):
method set_fps_limit (line 124) | def set_fps_limit(self, fps_limit):
method should_close (line 127) | def should_close(self):
method skip_frame (line 130) | def skip_frame(self):
method skip_frames (line 133) | def skip_frames(self, num): # Do not update window for the next N frames.
method is_skipping_frames (line 136) | def is_skipping_frames(self):
method capture_next_frame (line 139) | def capture_next_frame(self):
method pop_captured_frame (line 142) | def pop_captured_frame(self):
method pop_drag_and_drop_paths (line 147) | def pop_drag_and_drop_paths(self):
method draw_frame (line 152) | def draw_frame(self): # To be overridden by subclass.
method make_context_current (line 157) | def make_context_current(self):
method begin_frame (line 161) | def begin_frame(self):
method end_frame (line 198) | def end_frame(self):
method _attach_glfw_callbacks (line 218) | def _attach_glfw_callbacks(self):
method _glfw_key_callback (line 222) | def _glfw_key_callback(self, _window, key, _scancode, action, _mods):
method _glfw_drop_callback (line 226) | def _glfw_drop_callback(self, _window, paths):
FILE: gui_utils/imgui_utils.py
function set_default_style (line 14) | def set_default_style(color_scheme='dark', spacing=9, indent=23, scrollb...
function grayed_out (line 42) | def grayed_out(cond=True):
function item_width (line 70) | def item_width(width=None):
function scoped_by_object_id (line 80) | def scoped_by_object_id(method):
function button (line 90) | def button(label, width=0, enabled=True):
function collapsing_header (line 98) | def collapsing_header(text, visible=None, flags=0, default=False, enable...
function popup_button (line 112) | def popup_button(label, width=0, enabled=True):
function input_text (line 120) | def input_text(label, value, buffer_length, flags, width=None, help_text...
function drag_previous_control (line 137) | def drag_previous_control(enabled=True):
function drag_button (line 151) | def drag_button(label, width=0, enabled=True):
function drag_hidden_window (line 158) | def drag_hidden_window(label, x, y, width, height, enabled=True):
function click_hidden_window (line 171) | def click_hidden_window(label, x, y, width, height, img_w, img_h, enable...
FILE: gui_utils/imgui_window.py
class ImguiWindow (line 19) | class ImguiWindow(glfw_window.GlfwWindow):
method __init__ (line 20) | def __init__(self, *, title='ImguiWindow', font=None, font_sizes=range...
method close (line 45) | def close(self):
method _glfw_key_callback (line 56) | def _glfw_key_callback(self, *args):
method font_size (line 61) | def font_size(self):
method spacing (line 65) | def spacing(self):
method set_font_size (line 68) | def set_font_size(self, target): # Applied on next frame.
method begin_frame (line 71) | def begin_frame(self):
method end_frame (line 85) | def end_frame(self):
class _GlfwRenderer (line 95) | class _GlfwRenderer(imgui.integrations.glfw.GlfwRenderer):
method __init__ (line 96) | def __init__(self, *args, **kwargs):
method scroll_callback (line 100) | def scroll_callback(self, window, x_offset, y_offset):
FILE: gui_utils/text_utils.py
function get_default_font (line 22) | def get_default_font():
function get_pil_font (line 29) | def get_pil_font(font=None, size=32):
function get_array (line 36) | def get_array(string, *, dropshadow_radius: int=None, **kwargs):
function _get_array_priv (line 45) | def _get_array_priv(
function _get_array_impl (line 74) | def _get_array_impl(string, *, font=None, size=32, outline=0, outline_pa...
function _get_array_impl_dropshadow (line 96) | def _get_array_impl_dropshadow(string, *, font=None, size=32, radius: in...
function get_texture (line 120) | def get_texture(string, bilinear=True, mipmap=True, **kwargs):
FILE: legacy.py
function load_network_pkl (line 22) | def load_network_pkl(f, force_fp16=False):
class _TFNetworkStub (line 62) | class _TFNetworkStub(dnnlib.EasyDict):
class _LegacyUnpickler (line 65) | class _LegacyUnpickler(pickle.Unpickler):
method find_class (line 66) | def find_class(self, module, name):
function _collect_tf_params (line 73) | def _collect_tf_params(tf_net):
function _populate_module_params (line 86) | def _populate_module_params(module, *patterns):
function convert_tf_generator (line 107) | def convert_tf_generator(tf_G):
function convert_tf_discriminator (line 208) | def convert_tf_discriminator(tf_D):
function convert_network_pickle (line 297) | def convert_network_pickle(source, dest, force_fp16):
FILE: stylegan_human/PP_HumanSeg/deploy/infer.py
class DeployConfig (line 36) | class DeployConfig:
method __init__ (line 37) | def __init__(self, path):
method transforms (line 46) | def transforms(self):
method model (line 50) | def model(self):
method params (line 54) | def params(self):
method _load_transforms (line 57) | def _load_transforms(self, t_list):
class Predictor (line 67) | class Predictor:
method __init__ (line 68) | def __init__(self, args):
method preprocess (line 89) | def preprocess(self, img):
method run (line 97) | def run(self, img, bg):
method postprocess (line 117) | def postprocess(self, pred, img, ori_shape, bg):
FILE: stylegan_human/alignment.py
function angle_between_points (line 25) | def angle_between_points(p0,p1,p2):
function crop_img_with_padding (line 36) | def crop_img_with_padding(img, keypoints, rect):
function run (line 101) | def run(args):
FILE: stylegan_human/bg_white.py
function bg_white (line 8) | def bg_white(seg, raw, blur_level=3, gaussian=81):
function main (line 40) | def main(
FILE: stylegan_human/dnnlib/tflib/autosummary.py
function _create_var (line 47) | def _create_var(name: str, value_expr: TfExpression) -> TfExpression:
function autosummary (line 79) | def autosummary(name: str, value: TfExpressionEx, passthru: TfExpression...
function finalize_autosummaries (line 120) | def finalize_autosummaries() -> None:
function save_summaries (line 179) | def save_summaries(file_writer, global_step=None):
FILE: stylegan_human/dnnlib/tflib/custom_ops.py
function _find_compiler_bindir (line 38) | def _find_compiler_bindir():
function _get_compute_cap (line 44) | def _get_compute_cap(device):
function _get_cuda_gpu_arch_string (line 51) | def _get_cuda_gpu_arch_string():
function _run_cmd (line 58) | def _run_cmd(cmd):
function _prepare_nvcc_cli (line 65) | def _prepare_nvcc_cli(opts):
function get_plugin (line 89) | def get_plugin(cuda_file):
FILE: stylegan_human/dnnlib/tflib/network.py
function import_handler (line 31) | def import_handler(handler_func):
class Network (line 37) | class Network:
method __init__ (line 75) | def __init__(self, name: str = None, func_name: Any = None, **static_k...
method _init_fields (line 102) | def _init_fields(self) -> None:
method _init_graph (line 127) | def _init_graph(self) -> None:
method reset_own_vars (line 189) | def reset_own_vars(self) -> None:
method reset_vars (line 193) | def reset_vars(self) -> None:
method reset_trainables (line 197) | def reset_trainables(self) -> None:
method get_output_for (line 201) | def get_output_for(self, *in_expr: TfExpression, return_as_list: bool ...
method get_var_local_name (line 236) | def get_var_local_name(self, var_or_global_name: Union[TfExpression, s...
method find_var (line 242) | def find_var(self, var_or_local_name: Union[TfExpression, str]) -> TfE...
method get_var (line 247) | def get_var(self, var_or_local_name: Union[TfExpression, str]) -> np.n...
method set_var (line 252) | def set_var(self, var_or_local_name: Union[TfExpression, str], new_val...
method __getstate__ (line 257) | def __getstate__(self) -> dict:
method __setstate__ (line 269) | def __setstate__(self, state: dict) -> None:
method clone (line 303) | def clone(self, name: str = None, **new_static_kwargs) -> "Network":
method copy_own_vars_from (line 318) | def copy_own_vars_from(self, src_net: "Network") -> None:
method copy_vars_from (line 323) | def copy_vars_from(self, src_net: "Network") -> None:
method copy_trainables_from (line 328) | def copy_trainables_from(self, src_net: "Network") -> None:
method convert (line 333) | def convert(self, new_func_name: str, new_name: str = None, **new_stat...
method setup_as_moving_average_of (line 343) | def setup_as_moving_average_of(self, src_net: "Network", beta: TfExpre...
method run (line 355) | def run(self,
method list_ops (line 457) | def list_ops(self) -> List[TfExpression]:
method list_layers (line 465) | def list_layers(self) -> List[Tuple[str, TfExpression, List[TfExpressi...
method print_layers (line 508) | def print_layers(self, title: str = None, hide_layers_with_no_params: ...
method setup_weight_histograms (line 537) | def setup_weight_histograms(self, title: str = None) -> None:
function _handle_legacy_output_transforms (line 557) | def _handle_legacy_output_transforms(output_transform, dynamic_kwargs):
function _legacy_output_transform_func (line 577) | def _legacy_output_transform_func(*expr, out_mul=1.0, out_add=0.0, out_s...
FILE: stylegan_human/dnnlib/tflib/ops/fused_bias_act.py
function _get_plugin (line 17) | def _get_plugin():
function fused_bias_act (line 36) | def fused_bias_act(x, b=None, axis=1, act='linear', alpha=None, gain=Non...
function _fused_bias_act_ref (line 74) | def _fused_bias_act_ref(x, b, axis, act, alpha, gain):
function _fused_bias_act_cuda (line 102) | def _fused_bias_act_cuda(x, b, axis, act, alpha, gain):
FILE: stylegan_human/dnnlib/tflib/ops/upfirdn_2d.py
function _get_plugin (line 16) | def _get_plugin():
function upfirdn_2d (line 21) | def upfirdn_2d(x, k, upx=1, upy=1, downx=1, downy=1, padx0=0, padx1=0, p...
function _upfirdn_2d_ref (line 68) | def _upfirdn_2d_ref(x, k, upx, upy, downx, downy, padx0, padx1, pady0, p...
function _upfirdn_2d_cuda (line 107) | def _upfirdn_2d_cuda(x, k, upx, upy, downx, downy, padx0, padx1, pady0, ...
function filter_2d (line 146) | def filter_2d(x, k, gain=1, data_format='NCHW', impl='cuda'):
function upsample_2d (line 171) | def upsample_2d(x, k=None, factor=2, gain=1, data_format='NCHW', impl='c...
function downsample_2d (line 204) | def downsample_2d(x, k=None, factor=2, gain=1, data_format='NCHW', impl=...
function upsample_conv_2d (line 236) | def upsample_conv_2d(x, w, k=None, factor=2, gain=1, data_format='NCHW',...
function conv_downsample_2d (line 298) | def conv_downsample_2d(x, w, k=None, factor=2, gain=1, data_format='NCHW...
function _shape (line 339) | def _shape(tf_expr, dim_idx):
function _setup_kernel (line 346) | def _setup_kernel(k):
function _simple_upfirdn_2d (line 355) | def _simple_upfirdn_2d(x, k, up=1, down=1, pad0=0, pad1=0, data_format='...
FILE: stylegan_human/dnnlib/tflib/optimizer.py
class Optimizer (line 30) | class Optimizer:
method __init__ (line 42) | def __init__(self,
method _get_device (line 86) | def _get_device(self, device_name: str):
method register_gradients (line 116) | def register_gradients(self, loss: TfExpression, trainable_vars: Union...
method apply_updates (line 158) | def apply_updates(self, allow_no_op: bool = False) -> tf.Operation:
method reset_optimizer_state (line 268) | def reset_optimizer_state(self) -> None:
method get_loss_scaling_var (line 273) | def get_loss_scaling_var(self, device: str) -> Union[tf.Variable, None]:
method apply_loss_scaling (line 277) | def apply_loss_scaling(self, value: TfExpression) -> TfExpression:
method undo_loss_scaling (line 284) | def undo_loss_scaling(self, value: TfExpression) -> TfExpression:
class SimpleAdam (line 292) | class SimpleAdam:
method __init__ (line 295) | def __init__(self, name="Adam", learning_rate=0.001, beta1=0.9, beta2=...
method variables (line 303) | def variables(self):
method compute_gradients (line 306) | def compute_gradients(self, loss, var_list, gate_gradients=tf.train.Op...
method apply_gradients (line 310) | def apply_gradients(self, grads_and_vars):
FILE: stylegan_human/dnnlib/tflib/tfutil.py
function run (line 30) | def run(*args, **kwargs) -> Any:
function is_tf_expression (line 36) | def is_tf_expression(x: Any) -> bool:
function shape_to_list (line 41) | def shape_to_list(shape: Iterable[tf.Dimension]) -> List[Union[int, None]]:
function flatten (line 46) | def flatten(x: TfExpressionEx) -> TfExpression:
function log2 (line 52) | def log2(x: TfExpressionEx) -> TfExpression:
function exp2 (line 58) | def exp2(x: TfExpressionEx) -> TfExpression:
function lerp (line 64) | def lerp(a: TfExpressionEx, b: TfExpressionEx, t: TfExpressionEx) -> TfE...
function lerp_clip (line 70) | def lerp_clip(a: TfExpressionEx, b: TfExpressionEx, t: TfExpressionEx) -...
function absolute_name_scope (line 76) | def absolute_name_scope(scope: str) -> tf.name_scope:
function absolute_variable_scope (line 81) | def absolute_variable_scope(scope: str, **kwargs) -> tf.variable_scope:
function _sanitize_tf_config (line 86) | def _sanitize_tf_config(config_dict: dict = None) -> dict:
function init_tf (line 109) | def init_tf(config_dict: dict = None) -> None:
function assert_tf_initialized (line 137) | def assert_tf_initialized():
function create_session (line 143) | def create_session(config_dict: dict = None, force_as_default: bool = Fa...
function init_uninitialized_vars (line 166) | def init_uninitialized_vars(target_vars: List[tf.Variable] = None) -> None:
function set_vars (line 196) | def set_vars(var_to_value_dict: dict) -> None:
function create_var_with_large_initial_value (line 222) | def create_var_with_large_initial_value(initial_value: np.ndarray, *args...
function convert_images_from_uint8 (line 232) | def convert_images_from_uint8(images, drange=[-1,1], nhwc_to_nchw=False):
function convert_images_to_uint8 (line 242) | def convert_images_to_uint8(images, drange=[-1,1], nchw_to_nhwc=False, s...
FILE: stylegan_human/dnnlib/util.py
class EasyDict (line 41) | class EasyDict(dict):
method __getattr__ (line 44) | def __getattr__(self, name: str) -> Any:
method __setattr__ (line 50) | def __setattr__(self, name: str, value: Any) -> None:
method __delattr__ (line 53) | def __delattr__(self, name: str) -> None:
class Logger (line 57) | class Logger(object):
method __init__ (line 60) | def __init__(self, file_name: str = None, file_mode: str = "w", should...
method __enter__ (line 73) | def __enter__(self) -> "Logger":
method __exit__ (line 76) | def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> N...
method write (line 79) | def write(self, text: Union[str, bytes]) -> None:
method flush (line 94) | def flush(self) -> None:
method close (line 101) | def close(self) -> None:
function set_cache_dir (line 121) | def set_cache_dir(path: str) -> None:
function make_cache_dir_path (line 125) | def make_cache_dir_path(*paths: str) -> str:
function format_time (line 140) | def format_time(seconds: Union[int, float]) -> str:
function ask_yes_no (line 154) | def ask_yes_no(question: str) -> bool:
function tuple_product (line 164) | def tuple_product(t: Tuple) -> Any:
function get_dtype_and_ctype (line 188) | def get_dtype_and_ctype(type_obj: Any) -> Tuple[np.dtype, Any]:
function is_pickleable (line 211) | def is_pickleable(obj: Any) -> bool:
function get_module_from_obj_name (line 223) | def get_module_from_obj_name(obj_name: str) -> Tuple[types.ModuleType, s...
function get_obj_from_module (line 264) | def get_obj_from_module(module: types.ModuleType, obj_name: str) -> Any:
function get_obj_by_name (line 274) | def get_obj_by_name(name: str) -> Any:
function call_func_by_name (line 280) | def call_func_by_name(*args, func_name: str = None, **kwargs) -> Any:
function construct_class_by_name (line 289) | def construct_class_by_name(*args, class_name: str = None, **kwargs) -> ...
function get_module_dir_by_obj_name (line 294) | def get_module_dir_by_obj_name(obj_name: str) -> str:
function is_top_level_function (line 300) | def is_top_level_function(obj: Any) -> bool:
function get_top_level_function_name (line 305) | def get_top_level_function_name(obj: Any) -> str:
function list_dir_recursively_with_ignore (line 317) | def list_dir_recursively_with_ignore(dir_path: str, ignores: List[str] =...
function copy_files_and_create_dirs (line 350) | def copy_files_and_create_dirs(files: List[Tuple[str, str]]) -> None:
function is_url (line 366) | def is_url(obj: Any, allow_file_urls: bool = False) -> bool:
function open_url (line 384) | def open_url(url: str, cache_dir: str = None, num_attempts: int = 10, ve...
FILE: stylegan_human/edit.py
function main (line 55) | def main(
FILE: stylegan_human/edit/edit_helper.py
function conv_warper (line 10) | def conv_warper(layer, input, style, noise):
function decoder (line 58) | def decoder(G, style_space, latent, noise):
function encoder_ifg (line 76) | def encoder_ifg(G, noise, attr_name, truncation=1, truncation_latent=None,
function encoder_ss (line 122) | def encoder_ss(G, noise, attr_name, truncation=1, truncation_latent=None,
function encoder_sefa (line 180) | def encoder_sefa(G, noise, attr_name, truncation=1, truncation_latent=None,
FILE: stylegan_human/generate.py
function generate_images (line 48) | def generate_images(
FILE: stylegan_human/insetgan.py
class InsetGAN (line 19) | class InsetGAN(torch.nn.Module):
method __init__ (line 20) | def __init__(self, stylebody_ckpt, styleface_ckpt):
method loss_coarse (line 59) | def loss_coarse(self, A_face, B, p1=500, p2=0.05):
method get_border_mask (line 67) | def get_border_mask(A, x, spec):
method get_body_mask (line 76) | def get_body_mask(A, crop, padding=4):
method loss_border (line 81) | def loss_border(self, A_face, B, p1=10000, p2=2, spec=None):
method loss_body (line 87) | def loss_body(self, A, B, crop, p1=9000, p2=0.1):
method loss_face (line 94) | def loss_face(self, A, B, crop, p1=5000, p2=1.75):
method loss_reg (line 100) | def loss_reg(self, w, w_mean, p1, w_plus_delta=None, p2=None):
method detect_face_dlib (line 104) | def detect_face_dlib(self, img):
method dual_optimizer (line 117) | def dual_optimizer(self,
function main (line 347) | def main(
FILE: stylegan_human/interpolation.py
function lerp (line 22) | def lerp(code1, code2, alpha):
function slerp (line 27) | def slerp(code1, code2, alpha, DOT_THRESHOLD=0.9995): # Spherical linear...
function generate_image_from_z (line 49) | def generate_image_from_z(G, z, noise_mode, truncation_psi, device):
function get_concat_h (line 58) | def get_concat_h(im1, im2):
function make_latent_interp_animation (line 65) | def make_latent_interp_animation(G, code1, code2, img1, img2, num_interp...
function main (line 106) | def main(
FILE: stylegan_human/legacy.py
function load_network_pkl (line 22) | def load_network_pkl(f, force_fp16=False, G_only=False):
class _TFNetworkStub (line 75) | class _TFNetworkStub(dnnlib.EasyDict):
class _LegacyUnpickler (line 78) | class _LegacyUnpickler(pickle.Unpickler):
method find_class (line 79) | def find_class(self, module, name):
function num_range (line 86) | def num_range(s: str) -> List[int]:
function load_pkl (line 100) | def load_pkl(file_or_url):
function visual (line 107) | def visual(output, out_path):
function save_obj (line 120) | def save_obj(obj, path):
function convert_to_rgb (line 128) | def convert_to_rgb(state_ros, state_nv, ros_name, nv_name):
function convert_conv (line 135) | def convert_conv(state_ros, state_nv, ros_name, nv_name):
function convert_blur_kernel (line 143) | def convert_blur_kernel(state_ros, state_nv, level):
function determine_config (line 150) | def determine_config(state_nv):
function convert (line 161) | def convert(network_pkl, output_file, G_only=False):
FILE: stylegan_human/openpose/src/body.py
class Body (line 14) | class Body(object):
method __init__ (line 15) | def __init__(self, model_path):
method __call__ (line 23) | def __call__(self, oriImg):
FILE: stylegan_human/openpose/src/model.py
function make_layers (line 7) | def make_layers(block, no_relu_layers):
class bodypose_model (line 24) | class bodypose_model(nn.Module):
method __init__ (line 25) | def __init__(self):
method forward (line 114) | def forward(self, x):
class handpose_model (line 143) | class handpose_model(nn.Module):
method __init__ (line 144) | def __init__(self):
method forward (line 204) | def forward(self, x):
FILE: stylegan_human/openpose/src/util.py
function padRightDownCorner (line 12) | def padRightDownCorner(img, stride, padValue):
function transfer (line 35) | def transfer(model, model_weights):
function draw_bodypose (line 42) | def draw_bodypose(canvas, candidate, subset,show_number=False):
function npmax (line 81) | def npmax(array):
function npmax_with_score (line 89) | def npmax_with_score(array):
FILE: stylegan_human/pti/pti_models/e4e/encoders/helpers.py
class Flatten (line 11) | class Flatten(Module):
method forward (line 12) | def forward(self, input):
function l2_norm (line 16) | def l2_norm(input, axis=1):
class Bottleneck (line 22) | class Bottleneck(namedtuple('Block', ['in_channel', 'depth', 'stride'])):
function get_block (line 26) | def get_block(in_channel, depth, num_units, stride=2):
function get_blocks (line 30) | def get_blocks(num_layers):
class SEModule (line 57) | class SEModule(Module):
method __init__ (line 58) | def __init__(self, channels, reduction):
method forward (line 66) | def forward(self, x):
class bottleneck_IR (line 76) | class bottleneck_IR(Module):
method __init__ (line 77) | def __init__(self, in_channel, depth, stride):
method forward (line 92) | def forward(self, x):
class bottleneck_IR_SE (line 98) | class bottleneck_IR_SE(Module):
method __init__ (line 99) | def __init__(self, in_channel, depth, stride):
method forward (line 117) | def forward(self, x):
function _upsample_add (line 123) | def _upsample_add(x, y):
FILE: stylegan_human/pti/pti_models/e4e/encoders/model_irse.py
class Backbone (line 9) | class Backbone(Module):
method __init__ (line 10) | def __init__(self, input_size, num_layers, mode='ir', drop_ratio=0.4, ...
method forward (line 44) | def forward(self, x):
function IR_50 (line 51) | def IR_50(input_size):
function IR_101 (line 57) | def IR_101(input_size):
function IR_152 (line 63) | def IR_152(input_size):
function IR_SE_50 (line 69) | def IR_SE_50(input_size):
function IR_SE_101 (line 75) | def IR_SE_101(input_size):
function IR_SE_152 (line 81) | def IR_SE_152(input_size):
FILE: stylegan_human/pti/pti_models/e4e/encoders/psp_encoders.py
class ProgressiveStage (line 12) | class ProgressiveStage(Enum):
class GradualStyleBlock (line 34) | class GradualStyleBlock(Module):
method __init__ (line 35) | def __init__(self, in_c, out_c, spatial):
method forward (line 51) | def forward(self, x):
class GradualStyleEncoder (line 58) | class GradualStyleEncoder(Module):
method __init__ (line 59) | def __init__(self, num_layers, mode='ir', opts=None):
method forward (line 95) | def forward(self, x):
class Encoder4Editing (line 124) | class Encoder4Editing(Module):
method __init__ (line 125) | def __init__(self, num_layers, mode='ir', opts=None):
method get_deltas_starting_dimensions (line 165) | def get_deltas_starting_dimensions(self):
method set_progressive_stage (line 169) | def set_progressive_stage(self, new_stage: ProgressiveStage):
method forward (line 173) | def forward(self, x):
FILE: stylegan_human/pti/pti_models/e4e/latent_codes_pool.py
class LatentCodesPool (line 5) | class LatentCodesPool:
method __init__ (line 11) | def __init__(self, pool_size):
method query (line 21) | def query(self, ws):
method handle_w (line 42) | def handle_w(self, w, return_ws):
FILE: stylegan_human/pti/pti_models/e4e/psp.py
function get_keys (line 10) | def get_keys(d, name):
class pSp (line 17) | class pSp(nn.Module):
method __init__ (line 19) | def __init__(self, opts):
method set_encoder (line 29) | def set_encoder(self):
method load_weights (line 40) | def load_weights(self):
method forward (line 56) | def forward(self, x, resize=True, latent_mask=None, input_code=False, ...
method __load_latent_avg (line 93) | def __load_latent_avg(self, ckpt, repeat=None):
FILE: stylegan_human/pti/pti_models/e4e/stylegan2/model.py
class PixelNorm (line 11) | class PixelNorm(nn.Module):
method __init__ (line 12) | def __init__(self):
method forward (line 15) | def forward(self, input):
function make_kernel (line 19) | def make_kernel(k):
class Upsample (line 30) | class Upsample(nn.Module):
method __init__ (line 31) | def __init__(self, kernel, factor=2):
method forward (line 45) | def forward(self, input):
class Downsample (line 51) | class Downsample(nn.Module):
method __init__ (line 52) | def __init__(self, kernel, factor=2):
method forward (line 66) | def forward(self, input):
class Blur (line 72) | class Blur(nn.Module):
method __init__ (line 73) | def __init__(self, kernel, pad, upsample_factor=1):
method forward (line 85) | def forward(self, input):
class EqualConv2d (line 91) | class EqualConv2d(nn.Module):
method __init__ (line 92) | def __init__(
method forward (line 111) | def forward(self, input):
method __repr__ (line 122) | def __repr__(self):
class EqualLinear (line 129) | class EqualLinear(nn.Module):
method __init__ (line 130) | def __init__(
method forward (line 148) | def forward(self, input):
method __repr__ (line 160) | def __repr__(self):
class ScaledLeakyReLU (line 166) | class ScaledLeakyReLU(nn.Module):
method __init__ (line 167) | def __init__(self, negative_slope=0.2):
method forward (line 172) | def forward(self, input):
class ModulatedConv2d (line 178) | class ModulatedConv2d(nn.Module):
method __init__ (line 179) | def __init__(
method __repr__ (line 227) | def __repr__(self):
method forward (line 233) | def forward(self, input, style):
class NoiseInjection (line 277) | class NoiseInjection(nn.Module):
method __init__ (line 278) | def __init__(self):
method forward (line 283) | def forward(self, image, noise=None):
class ConstantInput (line 291) | class ConstantInput(nn.Module):
method __init__ (line 292) | def __init__(self, channel, size=4):
method forward (line 297) | def forward(self, input):
class StyledConv (line 304) | class StyledConv(nn.Module):
method __init__ (line 305) | def __init__(
method forward (line 332) | def forward(self, input, style, noise=None):
class ToRGB (line 341) | class ToRGB(nn.Module):
method __init__ (line 342) | def __init__(self, in_channel, style_dim, upsample=True, blur_kernel=[...
method forward (line 351) | def forward(self, input, style, skip=None):
class Generator (line 363) | class Generator(nn.Module):
method __init__ (line 364) | def __init__(
method make_noise (line 451) | def make_noise(self):
method mean_latent (line 462) | def mean_latent(self, n_latent):
method get_latent (line 470) | def get_latent(self, input):
method forward (line 473) | def forward(
class ConvLayer (line 552) | class ConvLayer(nn.Sequential):
method __init__ (line 553) | def __init__(
class ResBlock (line 601) | class ResBlock(nn.Module):
method __init__ (line 602) | def __init__(self, in_channel, out_channel, blur_kernel=[1, 3, 3, 1]):
method forward (line 612) | def forward(self, input):
class Discriminator (line 622) | class Discriminator(nn.Module):
method __init__ (line 623) | def __init__(self, size, channel_multiplier=2, blur_kernel=[1, 3, 3, 1]):
method forward (line 662) | def forward(self, input):
FILE: stylegan_human/pti/pti_models/e4e/stylegan2/op/fused_act.py
class FusedLeakyReLU (line 13) | class FusedLeakyReLU(nn.Module):
method __init__ (line 14) | def __init__(self, channel, negative_slope=0.2, scale=2 ** 0.5):
method forward (line 21) | def forward(self, input):
function fused_leaky_relu (line 25) | def fused_leaky_relu(input, bias, negative_slope=0.2, scale=2 ** 0.5):
FILE: stylegan_human/pti/pti_models/e4e/stylegan2/op/fused_bias_act.cpp
function fused_bias_act (line 11) | torch::Tensor fused_bias_act(const torch::Tensor& input, const torch::Te...
function PYBIND11_MODULE (line 19) | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
FILE: stylegan_human/pti/pti_models/e4e/stylegan2/op/upfirdn2d.cpp
function upfirdn2d (line 12) | torch::Tensor upfirdn2d(const torch::Tensor& input, const torch::Tensor&...
function PYBIND11_MODULE (line 21) | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
FILE: stylegan_human/pti/pti_models/e4e/stylegan2/op/upfirdn2d.py
function upfirdn2d (line 11) | def upfirdn2d(input, kernel, up=1, down=1, pad=(0, 0)):
function upfirdn2d_native (line 19) | def upfirdn2d_native(
FILE: stylegan_human/pti/training/coaches/base_coach.py
class BaseCoach (line 18) | class BaseCoach:
method __init__ (line 19) | def __init__(self, data_loader, use_wandb):
method restart_training (line 44) | def restart_training(self):
method get_inversion (line 55) | def get_inversion(self, w_path_dir, image_name, image):
method load_inversions (line 71) | def load_inversions(self, w_path_dir, image_name):
method calc_inversions (line 85) | def calc_inversions(self, image, image_name):
method train (line 98) | def train(self):
method configure_optimizers (line 101) | def configure_optimizers(self):
method calc_loss (line 106) | def calc_loss(self, generated_images, real_images, log_name, new_G, us...
method forward (line 127) | def forward(self, w):
method initilize_e4e (line 132) | def initilize_e4e(self):
method get_e4e_inversion (line 143) | def get_e4e_inversion(self, image):
FILE: stylegan_human/pti/training/coaches/localitly_regulizer.py
function l2_loss (line 8) | def l2_loss(real_images, generated_images):
class Space_Regulizer (line 13) | class Space_Regulizer:
method __init__ (line 14) | def __init__(self, original_G, lpips_net):
method get_morphed_w_code (line 19) | def get_morphed_w_code(self, new_w_code, fixed_w):
method get_image_from_ws (line 28) | def get_image_from_ws(self, w_codes, G):
method ball_holder_loss_lazy (line 31) | def ball_holder_loss_lazy(self, new_G, num_of_sampled_latents, w_batch...
method space_regulizer_loss (line 61) | def space_regulizer_loss(self, new_G, w_batch, use_wandb):
FILE: stylegan_human/pti/training/coaches/multi_id_coach.py
class MultiIDCoach (line 13) | class MultiIDCoach(BaseCoach):
method __init__ (line 15) | def __init__(self, data_loader, use_wandb):
method train (line 18) | def train(self):
FILE: stylegan_human/pti/training/coaches/single_id_coach.py
class SingleIDCoach (line 11) | class SingleIDCoach(BaseCoach):
method __init__ (line 13) | def __init__(self, data_loader, use_wandb):
method train (line 16) | def train(self):
FILE: stylegan_human/pti/training/projectors/w_plus_projector.py
function project (line 22) | def project(
FILE: stylegan_human/pti/training/projectors/w_projector.py
function project (line 22) | def project(
FILE: stylegan_human/run_pti.py
function run_PTI (line 16) | def run_PTI(run_name='', use_wandb=False, use_multi_id_training=False):
FILE: stylegan_human/style_mixing.py
function generate_style_mix (line 41) | def generate_style_mix(
FILE: stylegan_human/stylemixing_video.py
function style_mixing_video (line 49) | def style_mixing_video(network_pkl: str,
FILE: stylegan_human/torch_utils/custom_ops.py
function _find_compiler_bindir (line 32) | def _find_compiler_bindir():
function _get_mangled_gpu_name (line 45) | def _get_mangled_gpu_name():
function get_plugin (line 61) | def get_plugin(module_name, sources, **build_kwargs):
function get_plugin_v3 (line 142) | def get_plugin_v3(module_name, sources, headers=None, source_dir=None, *...
FILE: stylegan_human/torch_utils/misc.py
function constant (line 24) | def constant(value, shape=None, dtype=None, device=None, memory_format=N...
function nan_to_num (line 51) | def nan_to_num(input, nan=0.0, posinf=None, neginf=None, *, out=None): #...
class suppress_tracer_warnings (line 71) | class suppress_tracer_warnings(warnings.catch_warnings):
method __enter__ (line 72) | def __enter__(self):
function assert_shape (line 82) | def assert_shape(tensor, ref_shape):
function profiled_function (line 100) | def profiled_function(fn):
class InfiniteSampler (line 111) | class InfiniteSampler(torch.utils.data.Sampler):
method __init__ (line 112) | def __init__(self, dataset, rank=0, num_replicas=1, shuffle=True, seed...
method __iter__ (line 125) | def __iter__(self):
function params_and_buffers (line 147) | def params_and_buffers(module):
function named_params_and_buffers (line 151) | def named_params_and_buffers(module):
function copy_params_and_buffers (line 155) | def copy_params_and_buffers(src_module, dst_module, require_all=False):
function ddp_sync (line 169) | def ddp_sync(module, sync):
function check_ddp_consistency (line 180) | def check_ddp_consistency(module, ignore_regex=None):
function print_module_summary (line 194) | def print_module_summary(module, inputs, max_nesting=3, skip_redundant=T...
FILE: stylegan_human/torch_utils/models.py
class PixelNorm (line 19) | class PixelNorm(nn.Module):
method __init__ (line 20) | def __init__(self):
method forward (line 23) | def forward(self, input):
function make_kernel (line 27) | def make_kernel(k):
class Upsample (line 35) | class Upsample(nn.Module):
method __init__ (line 36) | def __init__(self, kernel, factor=2):
method forward (line 50) | def forward(self, input):
class Downsample (line 55) | class Downsample(nn.Module):
method __init__ (line 56) | def __init__(self, kernel, factor=2):
method forward (line 70) | def forward(self, input):
class Blur (line 75) | class Blur(nn.Module):
method __init__ (line 76) | def __init__(self, kernel, pad, upsample_factor=1):
method forward (line 88) | def forward(self, input):
class EqualConv2d (line 93) | class EqualConv2d(nn.Module):
method __init__ (line 94) | def __init__(
method forward (line 113) | def forward(self, input):
method __repr__ (line 123) | def __repr__(self):
class EqualLinear (line 130) | class EqualLinear(nn.Module):
method __init__ (line 131) | def __init__(
method forward (line 148) | def forward(self, input):
method __repr__ (line 158) | def __repr__(self):
class ScaledLeakyReLU (line 164) | class ScaledLeakyReLU(nn.Module):
method __init__ (line 165) | def __init__(self, negative_slope=0.2):
method forward (line 169) | def forward(self, input):
class ModulatedConv2d (line 174) | class ModulatedConv2d(nn.Module):
method __init__ (line 175) | def __init__(
method __repr__ (line 218) | def __repr__(self):
method forward (line 224) | def forward(self, input, style):
class NoiseInjection (line 268) | class NoiseInjection(nn.Module):
method __init__ (line 269) | def __init__(self):
method forward (line 273) | def forward(self, image, noise=None):
class ConstantInput (line 280) | class ConstantInput(nn.Module):
method __init__ (line 281) | def __init__(self, channel, size=4):
method forward (line 285) | def forward(self, input):
class StyledConv (line 291) | class StyledConv(nn.Module):
method __init__ (line 292) | def __init__(
method forward (line 315) | def forward(self, input, style, noise=None):
class ToRGB (line 322) | class ToRGB(nn.Module):
method __init__ (line 323) | def __init__(self, in_channel, style_dim, upsample=True, blur_kernel=[...
method forward (line 331) | def forward(self, input, style, skip=None):
class Generator (line 342) | class Generator(nn.Module):
method __init__ (line 343) | def __init__(
method make_noise (line 444) | def make_noise(self):
method mean_latent (line 455) | def mean_latent(self, n_latent):
method get_latent (line 463) | def get_latent(self, input):
method forward (line 466) | def forward(
class ConvLayer (line 558) | class ConvLayer(nn.Sequential):
method __init__ (line 559) | def __init__(
class ResBlock (line 606) | class ResBlock(nn.Module):
method __init__ (line 607) | def __init__(self, in_channel, out_channel, blur_kernel=[1, 3, 3, 1]):
method forward (line 617) | def forward(self, input):
class StyleDiscriminator (line 627) | class StyleDiscriminator(nn.Module):
method __init__ (line 628) | def __init__(
method forward (line 673) | def forward(self, input):
class StyleEncoder (line 701) | class StyleEncoder(nn.Module):
method __init__ (line 702) | def __init__(self, size, w_dim=512):
method forward (line 731) | def forward(self, input):
function kaiming_init (line 737) | def kaiming_init(m):
function normal_init (line 748) | def normal_init(m):
FILE: stylegan_human/torch_utils/models_face.py
class PixelNorm (line 17) | class PixelNorm(nn.Module):
method __init__ (line 18) | def __init__(self):
method forward (line 21) | def forward(self, input):
function make_kernel (line 25) | def make_kernel(k):
class Upsample (line 36) | class Upsample(nn.Module):
method __init__ (line 37) | def __init__(self, kernel, factor=2):
method forward (line 51) | def forward(self, input):
class Downsample (line 57) | class Downsample(nn.Module):
method __init__ (line 58) | def __init__(self, kernel, factor=2):
method forward (line 72) | def forward(self, input):
class Blur (line 78) | class Blur(nn.Module):
method __init__ (line 79) | def __init__(self, kernel, pad, upsample_factor=1):
method forward (line 91) | def forward(self, input):
class EqualConv2d (line 97) | class EqualConv2d(nn.Module):
method __init__ (line 98) | def __init__(
method forward (line 117) | def forward(self, input):
method __repr__ (line 128) | def __repr__(self):
class EqualLinear (line 135) | class EqualLinear(nn.Module):
method __init__ (line 136) | def __init__(
method forward (line 154) | def forward(self, input):
method __repr__ (line 166) | def __repr__(self):
class ScaledLeakyReLU (line 172) | class ScaledLeakyReLU(nn.Module):
method __init__ (line 173) | def __init__(self, negative_slope=0.2):
method forward (line 178) | def forward(self, input):
class ModulatedConv2d (line 184) | class ModulatedConv2d(nn.Module):
method __init__ (line 185) | def __init__(
method __repr__ (line 233) | def __repr__(self):
method forward (line 239) | def forward(self, input, style):
class NoiseInjection (line 283) | class NoiseInjection(nn.Module):
method __init__ (line 284) | def __init__(self):
method forward (line 289) | def forward(self, image, noise=None):
class ConstantInput (line 297) | class ConstantInput(nn.Module):
method __init__ (line 298) | def __init__(self, channel, size=4):
method forward (line 303) | def forward(self, input):
class StyledConv (line 310) | class StyledConv(nn.Module):
method __init__ (line 311) | def __init__(
method forward (line 338) | def forward(self, input, style, noise=None):
class ToRGB (line 347) | class ToRGB(nn.Module):
method __init__ (line 348) | def __init__(self, in_channel, style_dim, upsample=True, blur_kernel=[...
method forward (line 357) | def forward(self, input, style, skip=None):
class Generator (line 369) | class Generator(nn.Module):
method __init__ (line 370) | def __init__(
method make_noise (line 473) | def make_noise(self):
method mean_latent (line 484) | def mean_latent(self, n_latent):
method get_latent (line 492) | def get_latent(self, input):
method forward (line 495) | def forward(
class ConvLayer (line 585) | class ConvLayer(nn.Sequential):
method __init__ (line 586) | def __init__(
class ResBlock (line 634) | class ResBlock(nn.Module):
method __init__ (line 635) | def __init__(self, in_channel, out_channel, blur_kernel=[1, 3, 3, 1]):
method forward (line 645) | def forward(self, input):
class StyleDiscriminator (line 655) | class StyleDiscriminator(nn.Module):
method __init__ (line 656) | def __init__(
method forward (line 721) | def forward(self, input):
class StyleEncoder (line 749) | class StyleEncoder(nn.Module):
method __init__ (line 750) | def __init__(self, size, w_dim=512):
method forward (line 784) | def forward(self, input):
function kaiming_init (line 790) | def kaiming_init(m):
function normal_init (line 801) | def normal_init(m):
FILE: stylegan_human/torch_utils/op_edit/fused_act.py
class FusedLeakyReLUFunctionBackward (line 22) | class FusedLeakyReLUFunctionBackward(Function):
method forward (line 24) | def forward(ctx, grad_output, out, negative_slope, scale):
method backward (line 45) | def backward(ctx, gradgrad_input, gradgrad_bias):
class FusedLeakyReLUFunction (line 54) | class FusedLeakyReLUFunction(Function):
method forward (line 56) | def forward(ctx, input, bias, negative_slope, scale):
method backward (line 66) | def backward(ctx, grad_output):
class FusedLeakyReLU (line 76) | class FusedLeakyReLU(nn.Module):
method __init__ (line 77) | def __init__(self, channel, negative_slope=0.2, scale=2 ** 0.5):
method forward (line 84) | def forward(self, input):
function fused_leaky_relu (line 88) | def fused_leaky_relu(input, bias, negative_slope=0.2, scale=2 ** 0.5):
FILE: stylegan_human/torch_utils/op_edit/fused_bias_act.cpp
function fused_bias_act (line 13) | torch::Tensor fused_bias_act(const torch::Tensor& input, const torch::Te...
function PYBIND11_MODULE (line 21) | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
FILE: stylegan_human/torch_utils/op_edit/upfirdn2d.cpp
function upfirdn2d (line 14) | torch::Tensor upfirdn2d(const torch::Tensor& input, const torch::Tensor&...
function PYBIND11_MODULE (line 23) | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
FILE: stylegan_human/torch_utils/op_edit/upfirdn2d.py
class UpFirDn2dBackward (line 21) | class UpFirDn2dBackward(Function):
method forward (line 23) | def forward(
method backward (line 65) | def backward(ctx, gradgrad_input):
class UpFirDn2d (line 90) | class UpFirDn2d(Function):
method forward (line 92) | def forward(ctx, input, kernel, up, down, pad):
method backward (line 129) | def backward(ctx, grad_output):
function upfirdn2d (line 147) | def upfirdn2d(input, kernel, up=1, down=1, pad=(0, 0)):
function upfirdn2d_native (line 161) | def upfirdn2d_native(
FILE: stylegan_human/torch_utils/ops/bias_act.cpp
function has_same_layout (line 18) | static bool has_same_layout(torch::Tensor x, torch::Tensor y)
function bias_act (line 34) | static torch::Tensor bias_act(torch::Tensor x, torch::Tensor b, torch::T...
function PYBIND11_MODULE (line 96) | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
FILE: stylegan_human/torch_utils/ops/bias_act.h
type bias_act_kernel_params (line 14) | struct bias_act_kernel_params
FILE: stylegan_human/torch_utils/ops/bias_act.py
function _init (line 43) | def _init():
function bias_act (line 57) | def bias_act(x, b=None, dim=1, act='linear', alpha=None, gain=None, clam...
function _bias_act_ref (line 96) | def _bias_act_ref(x, b=None, dim=1, act='linear', alpha=None, gain=None,...
function _bias_act_cuda (line 131) | def _bias_act_cuda(dim=1, act='linear', alpha=None, gain=None, clamp=None):
FILE: stylegan_human/torch_utils/ops/conv2d_gradfix.py
function no_weight_gradients (line 28) | def no_weight_gradients():
function conv2d (line 37) | def conv2d(input, weight, bias=None, stride=1, padding=0, dilation=1, gr...
function conv_transpose2d (line 42) | def conv_transpose2d(input, weight, bias=None, stride=1, padding=0, outp...
function _should_use_custom_op (line 49) | def _should_use_custom_op(input):
function _tuple_of_ints (line 60) | def _tuple_of_ints(xs, ndim):
function _conv2d_gradfix (line 70) | def _conv2d_gradfix(transpose, weight_shape, stride, padding, output_pad...
FILE: stylegan_human/torch_utils/ops/conv2d_resample.py
function _get_weight_shape (line 23) | def _get_weight_shape(w):
function _conv2d_wrapper (line 31) | def _conv2d_wrapper(x, w, stride=1, padding=0, groups=1, transpose=False...
function conv2d_resample (line 61) | def conv2d_resample(x, w, f=None, up=1, down=1, padding=0, groups=1, fli...
FILE: stylegan_human/torch_utils/ops/filtered_lrelu.cpp
function filtered_lrelu (line 16) | static std::tuple<torch::Tensor, torch::Tensor, int> filtered_lrelu(
function filtered_lrelu_act (line 213) | static torch::Tensor filtered_lrelu_act(torch::Tensor x, torch::Tensor s...
function PYBIND11_MODULE (line 294) | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
FILE: stylegan_human/torch_utils/ops/filtered_lrelu.h
type filtered_lrelu_kernel_params (line 14) | struct filtered_lrelu_kernel_params
type filtered_lrelu_act_kernel_params (line 55) | struct filtered_lrelu_act_kernel_params
type filtered_lrelu_kernel_spec (line 73) | struct filtered_lrelu_kernel_spec
FILE: stylegan_human/torch_utils/ops/filtered_lrelu.py
function _init (line 23) | def _init():
function _get_filter_size (line 43) | def _get_filter_size(f):
function _parse_padding (line 50) | def _parse_padding(padding):
function filtered_lrelu (line 64) | def filtered_lrelu(x, fu=None, fd=None, b=None, up=1, down=1, padding=0,...
function _filtered_lrelu_ref (line 129) | def _filtered_lrelu_ref(x, fu=None, fd=None, b=None, up=1, down=1, paddi...
function _filtered_lrelu_cuda (line 167) | def _filtered_lrelu_cuda(up=1, down=1, padding=0, gain=np.sqrt(2), slope...
FILE: stylegan_human/torch_utils/ops/fma.py
function fma (line 17) | def fma(a, b, c): # => a * b + c
class _FusedMultiplyAdd (line 22) | class _FusedMultiplyAdd(torch.autograd.Function): # a * b + c
method forward (line 24) | def forward(ctx, a, b, c): # pylint: disable=arguments-differ
method backward (line 31) | def backward(ctx, dout): # pylint: disable=arguments-differ
function _unbroadcast (line 51) | def _unbroadcast(x, shape):
FILE: stylegan_human/torch_utils/ops/grid_sample_gradfix.py
function grid_sample (line 29) | def grid_sample(input, grid):
function _should_use_custom_op (line 36) | def _should_use_custom_op():
class _GridSample2dForward (line 46) | class _GridSample2dForward(torch.autograd.Function):
method forward (line 48) | def forward(ctx, input, grid):
method backward (line 56) | def backward(ctx, grad_output):
class _GridSample2dBackward (line 63) | class _GridSample2dBackward(torch.autograd.Function):
method forward (line 65) | def forward(ctx, grad_output, input, grid):
method backward (line 72) | def backward(ctx, grad2_grad_input, grad2_grad_grid):
FILE: stylegan_human/torch_utils/ops/upfirdn2d.cpp
function upfirdn2d (line 18) | static torch::Tensor upfirdn2d(torch::Tensor x, torch::Tensor f, int upx...
function PYBIND11_MODULE (line 100) | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
FILE: stylegan_human/torch_utils/ops/upfirdn2d.h
type upfirdn2d_kernel_params (line 16) | struct upfirdn2d_kernel_params
type upfirdn2d_kernel_spec (line 47) | struct upfirdn2d_kernel_spec
FILE: stylegan_human/torch_utils/ops/upfirdn2d.py
function _init (line 28) | def _init():
function _parse_scaling (line 39) | def _parse_scaling(scaling):
function _parse_padding (line 48) | def _parse_padding(padding):
function _get_filter_size (line 59) | def _get_filter_size(f):
function setup_filter (line 74) | def setup_filter(f, device=torch.device('cpu'), normalize=True, flip_fil...
function upfirdn2d (line 122) | def upfirdn2d(x, f, up=1, down=1, padding=0, flip_filter=False, gain=1, ...
function _upfirdn2d_ref (line 171) | def _upfirdn2d_ref(x, f, up=1, down=1, padding=0, flip_filter=False, gai...
function _upfirdn2d_cuda (line 216) | def _upfirdn2d_cuda(up=1, down=1, padding=0, flip_filter=False, gain=1):
function filter2d (line 274) | def filter2d(x, f, padding=0, flip_filter=False, gain=1, impl='cuda'):
function upsample2d (line 310) | def upsample2d(x, f, up=2, padding=0, flip_filter=False, gain=1, impl='c...
function downsample2d (line 349) | def downsample2d(x, f, down=2, padding=0, flip_filter=False, gain=1, imp...
FILE: stylegan_human/torch_utils/persistence.py
function persistent_class (line 37) | def persistent_class(orig_class):
function is_persistent (line 136) | def is_persistent(obj):
function import_hook (line 149) | def import_hook(hook):
function _reconstruct_persistent_obj (line 181) | def _reconstruct_persistent_obj(meta):
function _module_to_src (line 208) | def _module_to_src(module):
function _src_to_module (line 218) | def _src_to_module(src):
function _check_pickleable (line 233) | def _check_pickleable(obj):
FILE: stylegan_human/torch_utils/training_stats.py
function init_multiprocessing (line 36) | def init_multiprocessing(rank, sync_device):
function report (line 58) | def report(name, value):
function report0 (line 105) | def report0(name, value):
class Collector (line 115) | class Collector:
method __init__ (line 135) | def __init__(self, regex='.*', keep_previous=True):
method names (line 143) | def names(self):
method update (line 149) | def update(self):
method _get_delta (line 172) | def _get_delta(self, name):
method num (line 182) | def num(self, name):
method mean (line 190) | def mean(self, name):
method std (line 200) | def std(self, name):
method as_dict (line 214) | def as_dict(self):
method __getitem__ (line 228) | def __getitem__(self, name):
function _sync (line 236) | def _sync(names):
FILE: stylegan_human/training/augment.py
function matrix (line 48) | def matrix(*rows, device=None):
function translate2d (line 58) | def translate2d(tx, ty, **kwargs):
function translate3d (line 65) | def translate3d(tx, ty, tz, **kwargs):
function scale2d (line 73) | def scale2d(sx, sy, **kwargs):
function scale3d (line 80) | def scale3d(sx, sy, sz, **kwargs):
function rotate2d (line 88) | def rotate2d(theta, **kwargs):
function rotate3d (line 95) | def rotate3d(v, theta, **kwargs):
function translate2d_inv (line 105) | def translate2d_inv(tx, ty, **kwargs):
function scale2d_inv (line 108) | def scale2d_inv(sx, sy, **kwargs):
function rotate2d_inv (line 111) | def rotate2d_inv(theta, **kwargs):
class AugmentPipe (line 122) | class AugmentPipe(torch.nn.Module):
method __init__ (line 123) | def __init__(self,
method forward (line 186) | def forward(self, images, debug_percentile=None):
FILE: stylegan_human/training/dataset.py
class Dataset (line 26) | class Dataset(torch.utils.data.Dataset):
method __init__ (line 27) | def __init__(self,
method _get_raw_labels (line 53) | def _get_raw_labels(self):
method close (line 66) | def close(self): # to be overridden by subclass
method _load_raw_image (line 69) | def _load_raw_image(self, raw_idx): # to be overridden by subclass
method _load_raw_labels (line 72) | def _load_raw_labels(self): # to be overridden by subclass
method __getstate__ (line 75) | def __getstate__(self):
method __del__ (line 78) | def __del__(self):
method __len__ (line 84) | def __len__(self):
method __getitem__ (line 87) | def __getitem__(self, idx):
method get_label (line 97) | def get_label(self, idx):
method get_details (line 105) | def get_details(self, idx):
method name (line 113) | def name(self):
method image_shape (line 117) | def image_shape(self):
method num_channels (line 121) | def num_channels(self):
method resolution (line 126) | def resolution(self):
method label_shape (line 132) | def label_shape(self):
method label_dim (line 142) | def label_dim(self):
method has_labels (line 147) | def has_labels(self):
method has_onehot_labels (line 151) | def has_onehot_labels(self):
class ImageFolderDataset (line 156) | class ImageFolderDataset(Dataset):
method __init__ (line 157) | def __init__(self,
method _file_ext (line 186) | def _file_ext(fname):
method _get_zipfile (line 189) | def _get_zipfile(self):
method _open_file (line 195) | def _open_file(self, fname):
method close (line 202) | def close(self):
method __getstate__ (line 209) | def __getstate__(self):
method _load_raw_image (line 212) | def _load_raw_image(self, raw_idx):
method _load_raw_labels (line 224) | def _load_raw_labels(self):
FILE: stylegan_human/training/loss.py
class Loss (line 19) | class Loss:
method accumulate_gradients (line 20) | def accumulate_gradients(self, phase, real_img, real_c, gen_z, gen_c, ...
class StyleGAN2Loss (line 25) | class StyleGAN2Loss(Loss):
method __init__ (line 26) | def __init__(self, device, G, D, augment_pipe=None, r1_gamma=10, style...
method run_G (line 42) | def run_G(self, z, c, update_emas=False):
method run_D (line 52) | def run_D(self, img, c, blur_sigma=0, update_emas=False):
method accumulate_gradients (line 63) | def accumulate_gradients(self, phase, real_img, real_c, gen_z, gen_c, ...
FILE: stylegan_human/training/networks_stylegan2.py
function normalize_2nd_moment (line 27) | def normalize_2nd_moment(x, dim=1, eps=1e-8):
function modulated_conv2d (line 33) | def modulated_conv2d(
class FullyConnectedLayer (line 95) | class FullyConnectedLayer(torch.nn.Module):
method __init__ (line 96) | def __init__(self,
method forward (line 113) | def forward(self, x):
method extra_repr (line 128) | def extra_repr(self):
class Conv2dLayer (line 134) | class Conv2dLayer(torch.nn.Module):
method __init__ (line 135) | def __init__(self,
method forward (line 173) | def forward(self, x, gain=1):
method extra_repr (line 184) | def extra_repr(self):
class MappingNetwork (line 192) | class MappingNetwork(torch.nn.Module):
method __init__ (line 193) | def __init__(self,
method forward (line 232) | def forward(self, z, c, truncation_psi=1, truncation_cutoff=None, upda...
method extra_repr (line 269) | def extra_repr(self):
class SynthesisLayer (line 275) | class SynthesisLayer(torch.nn.Module):
method __init__ (line 276) | def __init__(self,
method forward (line 310) | def forward(self, x, w, noise_mode='random', fused_modconv=True, gain=1):
method extra_repr (line 331) | def extra_repr(self):
class ToRGBLayer (line 339) | class ToRGBLayer(torch.nn.Module):
method __init__ (line 340) | def __init__(self, in_channels, out_channels, w_dim, kernel_size=1, co...
method forward (line 352) | def forward(self, x, w, fused_modconv=True):
method extra_repr (line 358) | def extra_repr(self):
class SynthesisBlock (line 364) | class SynthesisBlock(torch.nn.Module):
method __init__ (line 365) | def __init__(self,
method forward (line 416) | def forward(self, x, img, ws, force_fp32=False, fused_modconv=None, up...
method extra_repr (line 462) | def extra_repr(self):
class SynthesisNetwork (line 468) | class SynthesisNetwork(torch.nn.Module):
method __init__ (line 469) | def __init__(self,
method forward (line 502) | def forward(self, ws, return_feature=False, **block_kwargs):
method extra_repr (line 524) | def extra_repr(self):
class Generator (line 533) | class Generator(torch.nn.Module):
method __init__ (line 534) | def __init__(self,
method forward (line 556) | def forward(self, z, c, truncation_psi=1, truncation_cutoff=None, upda...
function imresize (line 569) | def imresize(image, size):
class DiscriminatorBlock (line 585) | class DiscriminatorBlock(torch.nn.Module):
method __init__ (line 586) | def __init__(self,
method forward (line 636) | def forward(self, x, img, force_fp32=False):
method extra_repr (line 668) | def extra_repr(self):
class MinibatchStdLayer (line 674) | class MinibatchStdLayer(torch.nn.Module):
method __init__ (line 675) | def __init__(self, group_size, num_channels=1):
method forward (line 680) | def forward(self, x):
method extra_repr (line 697) | def extra_repr(self):
class DiscriminatorEpilogue (line 703) | class DiscriminatorEpilogue(torch.nn.Module):
method __init__ (line 704) | def __init__(self,
method forward (line 730) | def forward(self, x, img, cmap, force_fp32=False):
method extra_repr (line 758) | def extra_repr(self):
class Discriminator (line 764) | class Discriminator(torch.nn.Module):
method __init__ (line 765) | def __init__(self,
method forward (line 808) | def forward(self, img, c, update_emas=False, **block_kwargs):
method extra_repr (line 821) | def extra_repr(self):
FILE: stylegan_human/training/networks_stylegan3.py
function modulated_conv2d (line 26) | def modulated_conv2d(
class FullyConnectedLayer (line 70) | class FullyConnectedLayer(torch.nn.Module):
method __init__ (line 71) | def __init__(self,
method forward (line 90) | def forward(self, x):
method extra_repr (line 104) | def extra_repr(self):
class MappingNetwork (line 110) | class MappingNetwork(torch.nn.Module):
method __init__ (line 111) | def __init__(self,
method forward (line 136) | def forward(self, z, c, truncation_psi=1, truncation_cutoff=None, upda...
method extra_repr (line 164) | def extra_repr(self):
class SynthesisInput (line 170) | class SynthesisInput(torch.nn.Module):
method __init__ (line 171) | def __init__(self,
method forward (line 199) | def forward(self, w):
method extra_repr (line 246) | def extra_repr(self):
class SynthesisLayer (line 254) | class SynthesisLayer(torch.nn.Module):
method __init__ (line 255) | def __init__(self,
method forward (line 330) | def forward(self, x, w, noise_mode='random', force_fp32=False, update_...
method design_lowpass_filter (line 365) | def design_lowpass_filter(numtaps, cutoff, width, fs, radial=False):
method extra_repr (line 387) | def extra_repr(self):
class SynthesisNetwork (line 400) | class SynthesisNetwork(torch.nn.Module):
method __init__ (line 401) | def __init__(self,
method forward (line 465) | def forward(self, ws, **layer_kwargs):
method extra_repr (line 481) | def extra_repr(self):
class Generator (line 491) | class Generator(torch.nn.Module):
method __init__ (line 492) | def __init__(self,
method forward (line 513) | def forward(self, z, c, truncation_psi=1, truncation_cutoff=None, upda...
function imresize (line 527) | def imresize(image, size):
FILE: stylegan_human/training/training_loop.py
function setup_snapshot_image_grid (line 31) | def setup_snapshot_image_grid(training_set, random_seed=0):
function save_image_grid (line 70) | def save_image_grid(img, fname, drange, grid_size):
function training_loop (line 90) | def training_loop(
FILE: stylegan_human/training_scripts/sg2/train.py
class UserError (line 30) | class UserError(Exception):
function setup_training_loop_kwargs (line 35) | def setup_training_loop_kwargs(
function subprocess_fn (line 378) | def subprocess_fn(rank, args, temp_dir):
class CommaSeparatedList (line 403) | class CommaSeparatedList(click.ParamType):
method convert (line 406) | def convert(self, value, param, ctx):
function main (line 458) | def main(ctx, outdir, dry_run, **config_kwargs):
FILE: stylegan_human/training_scripts/sg2/training/dataset.py
class Dataset (line 29) | class Dataset(torch.utils.data.Dataset):
method __init__ (line 30) | def __init__(self,
method _get_raw_labels (line 59) | def _get_raw_labels(self):
method close (line 72) | def close(self): # to be overridden by subclass
method _load_raw_image (line 75) | def _load_raw_image(self, raw_idx): # to be overridden by subclass
method _load_raw_labels (line 78) | def _load_raw_labels(self): # to be overridden by subclass
method __getstate__ (line 81) | def __getstate__(self):
method __del__ (line 84) | def __del__(self):
method __len__ (line 90) | def __len__(self):
method __getitem__ (line 93) | def __getitem__(self, idx):
method get_label (line 103) | def get_label(self, idx):
method get_details (line 111) | def get_details(self, idx):
method name (line 119) | def name(self):
method image_shape (line 123) | def image_shape(self):
method num_channels (line 127) | def num_channels(self):
method resolution (line 132) | def resolution(self):
method label_shape (line 141) | def label_shape(self):
method label_dim (line 151) | def label_dim(self):
method has_labels (line 156) | def has_labels(self):
method has_onehot_labels (line 160) | def has_onehot_labels(self):
class ImageFolderDataset (line 165) | class ImageFolderDataset(Dataset):
method __init__ (line 166) | def __init__(self,
method _file_ext (line 204) | def _file_ext(fname):
method _get_zipfile (line 207) | def _get_zipfile(self):
method _open_file (line 213) | def _open_file(self, fname):
method close (line 220) | def close(self):
method __getstate__ (line 227) | def __getstate__(self):
method _load_raw_image (line 230) | def _load_raw_image(self, raw_idx): #load single image
method _load_raw_labels (line 242) | def _load_raw_labels(self):
FILE: stylegan_human/training_scripts/sg2/training/networks.py
function normalize_2nd_moment (line 23) | def normalize_2nd_moment(x, dim=1, eps=1e-8):
function modulated_conv2d (line 29) | def modulated_conv2d(
class FullyConnectedLayer (line 91) | class FullyConnectedLayer(torch.nn.Module):
method __init__ (line 92) | def __init__(self,
method forward (line 107) | def forward(self, x):
class Conv2dLayer (line 125) | class Conv2dLayer(torch.nn.Module):
method __init__ (line 126) | def __init__(self,
method forward (line 162) | def forward(self, x, gain=1):
class MappingNetwork (line 176) | class MappingNetwork(torch.nn.Module):
method __init__ (line 177) | def __init__(self,
method forward (line 216) | def forward(self, z, c, truncation_psi=1, truncation_cutoff=None, skip...
class SynthesisLayer (line 256) | class SynthesisLayer(torch.nn.Module):
method __init__ (line 257) | def __init__(self,
method forward (line 293) | def forward(self, x, w, noise_mode='random', fused_modconv=True, gain=1):
class ToRGBLayer (line 323) | class ToRGBLayer(torch.nn.Module):
method __init__ (line 324) | def __init__(self, in_channels, out_channels, w_dim, kernel_size=1, co...
method forward (line 333) | def forward(self, x, w, fused_modconv=True):
class SynthesisBlock (line 342) | class SynthesisBlock(torch.nn.Module):
method __init__ (line 343) | def __init__(self,
method forward (line 397) | def forward(self, x, img, ws, force_fp32=False, fused_modconv=None, **...
class SynthesisNetwork (line 448) | class SynthesisNetwork(torch.nn.Module):
method __init__ (line 449) | def __init__(self,
method forward (line 483) | def forward(self, ws, return_feature=False, **block_kwargs):
class Generator (line 508) | class Generator(torch.nn.Module):
method __init__ (line 509) | def __init__(self,
method forward (line 532) | def forward(self, z, c, truncation_psi=1, truncation_cutoff=None, inpu...
class DiscriminatorBlock (line 557) | class DiscriminatorBlock(torch.nn.Module):
method __init__ (line 558) | def __init__(self,
method forward (line 610) | def forward(self, x, img, force_fp32=False):
class MinibatchStdLayer (line 649) | class MinibatchStdLayer(torch.nn.Module):
method __init__ (line 650) | def __init__(self, group_size, num_channels=1):
method forward (line 655) | def forward(self, x):
class DiscriminatorEpilogue (line 675) | class DiscriminatorEpilogue(torch.nn.Module):
method __init__ (line 676) | def __init__(self,
method forward (line 709) | def forward(self, x, img, cmap, force_fp32=False):
class Discriminator (line 746) | class Discriminator(torch.nn.Module):
method __init__ (line 747) | def __init__(self,
method forward (line 792) | def forward(self, img, c, **block_kwargs):
FILE: stylegan_human/training_scripts/sg3/train.py
function subprocess_fn (line 29) | def subprocess_fn(rank, c, temp_dir):
function launch_training (line 53) | def launch_training(c, desc, outdir, dry_run):
function init_dataset_kwargs (line 104) | def init_dataset_kwargs(data, square=False):
function parse_comma_separated_list (line 120) | def parse_comma_separated_list(s):
function main (line 170) | def main(**kwargs):
FILE: stylegan_human/training_scripts/sg3/training/dataset.py
class Dataset (line 31) | class Dataset(torch.utils.data.Dataset):
method __init__ (line 32) | def __init__(self,
method _get_raw_labels (line 62) | def _get_raw_labels(self):
method close (line 75) | def close(self): # to be overridden by subclass
method _load_raw_image (line 78) | def _load_raw_image(self, raw_idx): # to be overridden by subclass
method _load_raw_labels (line 81) | def _load_raw_labels(self): # to be overridden by subclass
method __getstate__ (line 84) | def __getstate__(self):
method __del__ (line 87) | def __del__(self):
method __len__ (line 93) | def __len__(self):
method __getitem__ (line 96) | def __getitem__(self, idx):
method get_label (line 106) | def get_label(self, idx):
method get_details (line 114) | def get_details(self, idx):
method name (line 122) | def name(self):
method image_shape (line 126) | def image_shape(self):
method num_channels (line 130) | def num_channels(self):
method resolution (line 135) | def resolution(self):
method label_shape (line 144) | def label_shape(self):
method label_dim (line 154) | def label_dim(self):
method has_labels (line 159) | def has_labels(self):
method has_onehot_labels (line 163) | def has_onehot_labels(self):
class ImageFolderDataset (line 168) | class ImageFolderDataset(Dataset):
method __init__ (line 169) | def __init__(self,
method _file_ext (line 208) | def _file_ext(fname):
method _get_zipfile (line 211) | def _get_zipfile(self):
method _open_file (line 217) | def _open_file(self, fname):
method close (line 224) | def close(self):
method __getstate__ (line 231) | def __getstate__(self):
method _load_raw_image (line 234) | def _load_raw_image(self, raw_idx):
method _load_raw_labels (line 246) | def _load_raw_labels(self):
FILE: stylegan_human/training_scripts/sg3/training/networks_stylegan2.py
function normalize_2nd_moment (line 28) | def normalize_2nd_moment(x, dim=1, eps=1e-8):
function modulated_conv2d (line 34) | def modulated_conv2d(
class FullyConnectedLayer (line 96) | class FullyConnectedLayer(torch.nn.Module):
method __init__ (line 97) | def __init__(self,
method forward (line 114) | def forward(self, x):
method extra_repr (line 129) | def extra_repr(self):
class Conv2dLayer (line 135) | class Conv2dLayer(torch.nn.Module):
method __init__ (line 136) | def __init__(self,
method forward (line 174) | def forward(self, x, gain=1):
method extra_repr (line 185) | def extra_repr(self):
class MappingNetwork (line 193) | class MappingNetwork(torch.nn.Module):
method __init__ (line 194) | def __init__(self,
method forward (line 233) | def forward(self, z, c, truncation_psi=1, truncation_cutoff=None, upda...
method extra_repr (line 270) | def extra_repr(self):
class SynthesisLayer (line 276) | class SynthesisLayer(torch.nn.Module):
method __init__ (line 277) | def __init__(self,
method forward (line 316) | def forward(self, x, w, noise_mode='random', fused_modconv=True, gain=1):
method extra_repr (line 343) | def extra_repr(self):
class ToRGBLayer (line 351) | class ToRGBLayer(torch.nn.Module):
method __init__ (line 352) | def __init__(self, in_channels, out_channels, w_dim, kernel_size=1, co...
method forward (line 364) | def forward(self, x, w, fused_modconv=True):
method extra_repr (line 370) | def extra_repr(self):
class SynthesisBlock (line 376) | class SynthesisBlock(torch.nn.Module):
method __init__ (line 377) | def __init__(self,
method forward (line 433) | def forward(self, x, img, ws, force_fp32=False, fused_modconv=None, up...
method extra_repr (line 485) | def extra_repr(self):
class SynthesisNetwork (line 491) | class SynthesisNetwork(torch.nn.Module):
method __init__ (line 492) | def __init__(self,
method forward (line 527) | def forward(self, ws, **block_kwargs):
method extra_repr (line 544) | def extra_repr(self):
class Generator (line 553) | class Generator(torch.nn.Module):
method __init__ (line 554) | def __init__(self,
method forward (line 575) | def forward(self, z, c, truncation_psi=1, truncation_cutoff=None, upda...
class DiscriminatorBlock (line 583) | class DiscriminatorBlock(torch.nn.Module):
method __init__ (line 584) | def __init__(self,
method forward (line 636) | def forward(self, x, img, force_fp32=False):
method extra_repr (line 674) | def extra_repr(self):
class MinibatchStdLayer (line 680) | class MinibatchStdLayer(torch.nn.Module):
method __init__ (line 681) | def __init__(self, group_size, num_channels=1):
method forward (line 686) | def forward(self, x):
method extra_repr (line 703) | def extra_repr(self):
class DiscriminatorEpilogue (line 709) | class DiscriminatorEpilogue(torch.nn.Module):
method __init__ (line 710) | def __init__(self,
method forward (line 743) | def forward(self, x, img, cmap, force_fp32=False):
method extra_repr (line 779) | def extra_repr(self):
class Discriminator (line 785) | class Discriminator(torch.nn.Module):
method __init__ (line 786) | def __init__(self,
method forward (line 831) | def forward(self, img, c, update_emas=False, **block_kwargs):
method extra_repr (line 844) | def extra_repr(self):
FILE: stylegan_human/training_scripts/sg3/training/networks_stylegan3.py
function modulated_conv2d (line 27) | def modulated_conv2d(
class FullyConnectedLayer (line 71) | class FullyConnectedLayer(torch.nn.Module):
method __init__ (line 72) | def __init__(self,
method forward (line 91) | def forward(self, x):
method extra_repr (line 105) | def extra_repr(self):
class MappingNetwork (line 111) | class MappingNetwork(torch.nn.Module):
method __init__ (line 112) | def __init__(self,
method forward (line 137) | def forward(self, z, c, truncation_psi=1, truncation_cutoff=None, upda...
method extra_repr (line 165) | def extra_repr(self):
class SynthesisInput (line 171) | class SynthesisInput(torch.nn.Module):
method __init__ (line 172) | def __init__(self,
method forward (line 205) | def forward(self, w):
method extra_repr (line 252) | def extra_repr(self):
class SynthesisLayer (line 260) | class SynthesisLayer(torch.nn.Module):
method __init__ (line 261) | def __init__(self,
method forward (line 344) | def forward(self, x, w, noise_mode='random', force_fp32=False, update_...
method design_lowpass_filter (line 379) | def design_lowpass_filter(numtaps, cutoff, width, fs, radial=False):
method extra_repr (line 401) | def extra_repr(self):
class SynthesisNetwork (line 414) | class SynthesisNetwork(torch.nn.Module):
method __init__ (line 415) | def __init__(self,
method forward (line 483) | def forward(self, ws, **layer_kwargs):
method extra_repr (line 502) | def extra_repr(self):
class Generator (line 512) | class Generator(torch.nn.Module):
method __init__ (line 513) | def __init__(self,
method forward (line 534) | def forward(self, z, c, truncation_psi=1, truncation_cutoff=None, upda...
FILE: stylegan_human/utils/ImagesDataset.py
class ImagesDataset (line 11) | class ImagesDataset(Dataset):
method __init__ (line 13) | def __init__(self, source_root, source_transform=None):
method __len__ (line 17) | def __len__(self):
method __getitem__ (line 20) | def __getitem__(self, index):
FILE: stylegan_human/utils/data_utils.py
function is_image_file (line 14) | def is_image_file(filename):
function tensor2im (line 18) | def tensor2im(var):
function make_dataset (line 28) | def make_dataset(dir):
FILE: stylegan_human/utils/face_alignment.py
function get_landmark (line 12) | def get_landmark(img, detector, predictor):
function align_face_for_insetgan (line 34) | def align_face_for_insetgan(img, detector, predictor, output_size=256):
function align_face_for_projector (line 137) | def align_face_for_projector(img, detector, predictor, output_size):
function reverse_quad_transform (line 227) | def reverse_quad_transform(image, quad_to_map_to, alpha):
FILE: stylegan_human/utils/log_utils.py
function log_image_from_w (line 12) | def log_image_from_w(w, G, name):
function log_images_from_w (line 21) | def log_images_from_w(ws, G, names):
function plot_image_from_w (line 27) | def plot_image_from_w(w, G):
function plot_image (line 34) | def plot_image(img):
function save_image (line 41) | def save_image(name, method_type, results_dir, image, run_id):
function save_w (line 45) | def save_w(w, G, name, method_type, results_dir):
function save_concat_image (line 51) | def save_concat_image(base_dir, image_latents, new_inv_image_latent, new_G,
function save_single_image (line 65) | def save_single_image(base_dir, image_latent, G, file_name):
function create_alongside_images (line 71) | def create_alongside_images(images):
function get_image_from_w (line 76) | def get_image_from_w(w, G):
FILE: stylegan_human/utils/models_utils.py
function toogle_grad (line 10) | def toogle_grad(model, flag=True):
function load_tuned_G (line 15) | def load_tuned_G(run_id, type):
function load_old_G (line 24) | def load_old_G():
FILE: stylegan_human/utils/util.py
function visual (line 10) | def visual(output, out_path):
function get_lr (line 21) | def get_lr(t, initial_lr, rampdown=0.25, rampup=0.05):
function latent_noise (line 30) | def latent_noise(latent, strength):
function noise_regularize_ (line 35) | def noise_regularize_(noises):
function noise_normalize_ (line 58) | def noise_normalize_(noises):
function tensor_to_numpy (line 66) | def tensor_to_numpy(x):
function numpy_to_tensor (line 73) | def numpy_to_tensor(x):
function tensor_to_pil (line 79) | def tensor_to_pil(x):
FILE: torch_utils/custom_ops.py
function _find_compiler_bindir (line 29) | def _find_compiler_bindir():
function _get_mangled_gpu_name (line 44) | def _get_mangled_gpu_name():
function get_plugin (line 59) | def get_plugin(module_name, sources, headers=None, source_dir=None, **bu...
FILE: torch_utils/misc.py
function constant (line 22) | def constant(value, shape=None, dtype=None, device=None, memory_format=N...
function nan_to_num (line 49) | def nan_to_num(input, nan=0.0, posinf=None, neginf=None, *, out=None): #...
function suppress_tracer_warnings (line 71) | def suppress_tracer_warnings():
function assert_shape (line 82) | def assert_shape(tensor, ref_shape):
function profiled_function (line 100) | def profiled_function(fn):
class InfiniteSampler (line 111) | class InfiniteSampler(torch.utils.data.Sampler):
method __init__ (line 112) | def __init__(self, dataset, rank=0, num_replicas=1, shuffle=True, seed...
method __iter__ (line 125) | def __iter__(self):
function params_and_buffers (line 147) | def params_and_buffers(module):
function named_params_and_buffers (line 151) | def named_params_and_buffers(module):
function copy_params_and_buffers (line 155) | def copy_params_and_buffers(src_module, dst_module, require_all=False):
function ddp_sync (line 169) | def ddp_sync(module, sync):
function check_ddp_consistency (line 180) | def check_ddp_consistency(module, ignore_regex=None):
function print_module_summary (line 196) | def print_module_summary(module, inputs, max_nesting=3, skip_redundant=T...
FILE: torch_utils/ops/bias_act.cpp
function has_same_layout (line 16) | static bool has_same_layout(torch::Tensor x, torch::Tensor y)
function bias_act (line 32) | static torch::Tensor bias_act(torch::Tensor x, torch::Tensor b, torch::T...
function PYBIND11_MODULE (line 94) | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
FILE: torch_utils/ops/bias_act.h
type bias_act_kernel_params (line 12) | struct bias_act_kernel_params
FILE: torch_utils/ops/bias_act.py
function _init (line 38) | def _init():
function bias_act (line 52) | def bias_act(x, b=None, dim=1, act='linear', alpha=None, gain=None, clam...
function _bias_act_ref (line 91) | def _bias_act_ref(x, b=None, dim=1, act='linear', alpha=None, gain=None,...
function _bias_act_cuda (line 126) | def _bias_act_cuda(dim=1, act='linear', alpha=None, gain=None, clamp=None):
FILE: torch_utils/ops/conv2d_gradfix.py
function no_weight_gradients (line 25) | def no_weight_gradients(disable=True):
function conv2d (line 35) | def conv2d(input, weight, bias=None, stride=1, padding=0, dilation=1, gr...
function conv_transpose2d (line 40) | def conv_transpose2d(input, weight, bias=None, stride=1, padding=0, outp...
function _should_use_custom_op (line 47) | def _should_use_custom_op(input):
function _tuple_of_ints (line 55) | def _tuple_of_ints(xs, ndim):
function _conv2d_gradfix (line 66) | def _conv2d_gradfix(transpose, weight_shape, stride, padding, output_pad...
FILE: torch_utils/ops/conv2d_resample.py
function _get_weight_shape (line 21) | def _get_weight_shape(w):
function _conv2d_wrapper (line 29) | def _conv2d_wrapper(x, w, stride=1, padding=0, groups=1, transpose=False...
function conv2d_resample (line 46) | def conv2d_resample(x, w, f=None, up=1, down=1, padding=0, groups=1, fli...
FILE: torch_utils/ops/filtered_lrelu.cpp
function filtered_lrelu (line 16) | static std::tuple<torch::Tensor, torch::Tensor, int> filtered_lrelu(
function filtered_lrelu_act (line 213) | static torch::Tensor filtered_lrelu_act(torch::Tensor x, torch::Tensor s...
function PYBIND11_MODULE (line 294) | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
FILE: torch_utils/ops/filtered_lrelu.h
type filtered_lrelu_kernel_params (line 14) | struct filtered_lrelu_kernel_params
type filtered_lrelu_act_kernel_params (line 55) | struct filtered_lrelu_act_kernel_params
type filtered_lrelu_kernel_spec (line 73) | struct filtered_lrelu_kernel_spec
FILE: torch_utils/ops/filtered_lrelu.py
function _init (line 23) | def _init():
function _get_filter_size (line 35) | def _get_filter_size(f):
function _parse_padding (line 42) | def _parse_padding(padding):
function filtered_lrelu (line 56) | def filtered_lrelu(x, fu=None, fd=None, b=None, up=1, down=1, padding=0,...
function _filtered_lrelu_ref (line 121) | def _filtered_lrelu_ref(x, fu=None, fd=None, b=None, up=1, down=1, paddi...
function _filtered_lrelu_cuda (line 159) | def _filtered_lrelu_cuda(up=1, down=1, padding=0, gain=np.sqrt(2), slope...
FILE: torch_utils/ops/fma.py
function fma (line 15) | def fma(a, b, c): # => a * b + c
class _FusedMultiplyAdd (line 20) | class _FusedMultiplyAdd(torch.autograd.Function): # a * b + c
method forward (line 22) | def forward(ctx, a, b, c): # pylint: disable=arguments-differ
method backward (line 29) | def backward(ctx, dout): # pylint: disable=arguments-differ
function _unbroadcast (line 49) | def _unbroadcast(x, shape):
FILE: torch_utils/ops/grid_sample_gradfix.py
function grid_sample (line 26) | def grid_sample(input, grid):
function _should_use_custom_op (line 33) | def _should_use_custom_op():
class _GridSample2dForward (line 38) | class _GridSample2dForward(torch.autograd.Function):
method forward (line 40) | def forward(ctx, input, grid):
method backward (line 48) | def backward(ctx, grad_output):
class _GridSample2dBackward (line 55) | class _GridSample2dBackward(torch.autograd.Function):
method forward (line 57) | def forward(ctx, grad_output, input, grid):
method backward (line 64) | def backward(ctx, grad2_grad_input, grad2_grad_grid):
FILE: torch_utils/ops/upfirdn2d.cpp
function upfirdn2d (line 16) | static torch::Tensor upfirdn2d(torch::Tensor x, torch::Tensor f, int upx...
function PYBIND11_MODULE (line 102) | PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
FILE: torch_utils/ops/upfirdn2d.h
type upfirdn2d_kernel_params (line 14) | struct upfirdn2d_kernel_params
type upfirdn2d_kernel_spec (line 45) | struct upfirdn2d_kernel_spec
FILE: torch_utils/ops/upfirdn2d.py
function _init (line 23) | def _init():
function _parse_scaling (line 35) | def _parse_scaling(scaling):
function _parse_padding (line 44) | def _parse_padding(padding):
function _get_filter_size (line 55) | def _get_filter_size(f):
function setup_filter (line 70) | def setup_filter(f, device=torch.device('cpu'), normalize=True, flip_fil...
function upfirdn2d (line 118) | def upfirdn2d(x, f, up=1, down=1, padding=0, flip_filter=False, gain=1, ...
function _upfirdn2d_ref (line 167) | def _upfirdn2d_ref(x, f, up=1, down=1, padding=0, flip_filter=False, gai...
function _upfirdn2d_cuda (line 217) | def _upfirdn2d_cuda(up=1, down=1, padding=0, flip_filter=False, gain=1):
function filter2d (line 277) | def filter2d(x, f, padding=0, flip_filter=False, gain=1, impl='cuda'):
function upsample2d (line 313) | def upsample2d(x, f, up=2, padding=0, flip_filter=False, gain=1, impl='c...
function downsample2d (line 352) | def downsample2d(x, f, down=2, padding=0, flip_filter=False, gain=1, imp...
FILE: torch_utils/persistence.py
function persistent_class (line 35) | def persistent_class(orig_class):
function is_persistent (line 134) | def is_persistent(obj):
function import_hook (line 147) | def import_hook(hook):
function _reconstruct_persistent_obj (line 179) | def _reconstruct_persistent_obj(meta):
function _module_to_src (line 206) | def _module_to_src(module):
function _src_to_module (line 216) | def _src_to_module(src):
function _check_pickleable (line 231) | def _check_pickleable(obj):
FILE: torch_utils/training_stats.py
function init_multiprocessing (line 34) | def init_multiprocessing(rank, sync_device):
function report (line 56) | def report(name, value):
function report0 (line 103) | def report0(name, value):
class Collector (line 113) | class Collector:
method __init__ (line 133) | def __init__(self, regex='.*', keep_previous=True):
method names (line 141) | def names(self):
method update (line 147) | def update(self):
method _get_delta (line 170) | def _get_delta(self, name):
method num (line 180) | def num(self, name):
method mean (line 188) | def mean(self, name):
method std (line 198) | def std(self, name):
method as_dict (line 212) | def as_dict(self):
method __getitem__ (line 226) | def __getitem__(self, name):
function _sync (line 234) | def _sync(names):
FILE: training/augment.py
function matrix (line 48) | def matrix(*rows, device=None):
function translate2d (line 58) | def translate2d(tx, ty, **kwargs):
function translate3d (line 65) | def translate3d(tx, ty, tz, **kwargs):
function scale2d (line 73) | def scale2d(sx, sy, **kwargs):
function scale3d (line 80) | def scale3d(sx, sy, sz, **kwargs):
function rotate2d (line 88) | def rotate2d(theta, **kwargs):
function rotate3d (line 95) | def rotate3d(v, theta, **kwargs):
function translate2d_inv (line 105) | def translate2d_inv(tx, ty, **kwargs):
function scale2d_inv (line 108) | def scale2d_inv(sx, sy, **kwargs):
function rotate2d_inv (line 111) | def rotate2d_inv(theta, **kwargs):
class AugmentPipe (line 122) | class AugmentPipe(torch.nn.Module):
method __init__ (line 123) | def __init__(self,
method forward (line 186) | def forward(self, images, debug_percentile=None):
FILE: training/dataset.py
class Dataset (line 26) | class Dataset(torch.utils.data.Dataset):
method __init__ (line 27) | def __init__(self,
method _get_raw_labels (line 53) | def _get_raw_labels(self):
method close (line 66) | def close(self): # to be overridden by subclass
method _load_raw_image (line 69) | def _load_raw_image(self, raw_idx): # to be overridden by subclass
method _load_raw_labels (line 72) | def _load_raw_labels(self): # to be overridden by subclass
method __getstate__ (line 75) | def __getstate__(self):
method __del__ (line 78) | def __del__(self):
method __len__ (line 84) | def __len__(self):
method __getitem__ (line 87) | def __getitem__(self, idx):
method get_label (line 97) | def get_label(self, idx):
method get_details (line 105) | def get_details(self, idx):
method name (line 113) | def name(self):
method image_shape (line 117) | def image_shape(self):
method num_channels (line 121) | def num_channels(self):
method resolution (line 126) | def resolution(self):
method label_shape (line 132) | def label_shape(self):
method label_dim (line 142) | def label_dim(self):
method has_labels (line 147) | def has_labels(self):
method has_onehot_labels (line 151) | def has_onehot_labels(self):
class ImageFolderDataset (line 156) | class ImageFolderDataset(Dataset):
method __init__ (line 157) | def __init__(self,
method _file_ext (line 186) | def _file_ext(fname):
method _get_zipfile (line 189) | def _get_zipfile(self):
method _open_file (line 195) | def _open_file(self, fname):
method close (line 202) | def close(self):
method __getstate__ (line 209) | def __getstate__(self):
method _load_raw_image (line 212) | def _load_raw_image(self, raw_idx):
method _load_raw_labels (line 224) | def _load_raw_labels(self):
FILE: training/loss.py
class Loss (line 19) | class Loss:
method accumulate_gradients (line 20) | def accumulate_gradients(self, phase, real_img, real_c, gen_z, gen_c, ...
class StyleGAN2Loss (line 25) | class StyleGAN2Loss(Loss):
method __init__ (line 26) | def __init__(self, device, G, D, augment_pipe=None, r1_gamma=10, style...
method run_G (line 42) | def run_G(self, z, c, update_emas=False):
method run_D (line 52) | def run_D(self, img, c, blur_sigma=0, update_emas=False):
method accumulate_gradients (line 63) | def accumulate_gradients(self, phase, real_img, real_c, gen_z, gen_c, ...
FILE: training/networks_stylegan2.py
function normalize_2nd_moment (line 27) | def normalize_2nd_moment(x, dim=1, eps=1e-8):
function modulated_conv2d (line 33) | def modulated_conv2d(
class FullyConnectedLayer (line 95) | class FullyConnectedLayer(torch.nn.Module):
method __init__ (line 96) | def __init__(self,
method forward (line 113) | def forward(self, x):
method extra_repr (line 128) | def extra_repr(self):
class Conv2dLayer (line 134) | class Conv2dLayer(torch.nn.Module):
method __init__ (line 135) | def __init__(self,
method forward (line 173) | def forward(self, x, gain=1):
method extra_repr (line 184) | def extra_repr(self):
class MappingNetwork (line 192) | class MappingNetwork(torch.nn.Module):
method __init__ (line 193) | def __init__(self,
method forward (line 232) | def forward(self, z, c, truncation_psi=1, truncation_cutoff=None, upda...
method extra_repr (line 269) | def extra_repr(self):
class SynthesisLayer (line 275) | class SynthesisLayer(torch.nn.Module):
method __init__ (line 276) | def __init__(self,
method forward (line 310) | def forward(self, x, w, noise_mode='random', fused_modconv=True, gain=1):
method extra_repr (line 331) | def extra_repr(self):
class ToRGBLayer (line 339) | class ToRGBLayer(torch.nn.Module):
method __init__ (line 340) | def __init__(self, in_channels, out_channels, w_dim, kernel_size=1, co...
method forward (line 352) | def forward(self, x, w, fused_modconv=True):
method extra_repr (line 358) | def extra_repr(self):
class SynthesisBlock (line 364) | class SynthesisBlock(torch.nn.Module):
method __init__ (line 365) | def __init__(self,
method forward (line 416) | def forward(self, x, img, ws, force_fp32=False, fused_modconv=None, up...
method extra_repr (line 462) | def extra_repr(self):
class SynthesisNetwork (line 468) | class SynthesisNetwork(torch.nn.Module):
method __init__ (line 469) | def __init__(self,
method forward (line 502) | def forward(self, ws, return_feature=False, **block_kwargs):
method extra_repr (line 524) | def extra_repr(self):
class Generator (line 533) | class Generator(torch.nn.Module):
method __init__ (line 534) | def __init__(self,
method forward (line 558) | def forward(self, z, c, truncation_psi=1, truncation_cutoff=None, upda...
function imresize (line 576) | def imresize(image, size):
class DiscriminatorBlock (line 592) | class DiscriminatorBlock(torch.nn.Module):
method __init__ (line 593) | def __init__(self,
method forward (line 643) | def forward(self, x, img, force_fp32=False):
method extra_repr (line 675) | def extra_repr(self):
class MinibatchStdLayer (line 681) | class MinibatchStdLayer(torch.nn.Module):
method __init__ (line 682) | def __init__(self, group_size, num_channels=1):
method forward (line 687) | def forward(self, x):
method extra_repr (line 704) | def extra_repr(self):
class DiscriminatorEpilogue (line 710) | class DiscriminatorEpilogue(torch.nn.Module):
method __init__ (line 711) | def __init__(self,
method forward (line 737) | def forward(self, x, img, cmap, force_fp32=False):
method extra_repr (line 765) | def extra_repr(self):
class Discriminator (line 771) | class Discriminator(torch.nn.Module):
method __init__ (line 772) | def __init__(self,
method forward (line 815) | def forward(self, img, c, update_emas=False, **block_kwargs):
method extra_repr (line 828) | def extra_repr(self):
FILE: training/networks_stylegan3.py
function modulated_conv2d (line 26) | def modulated_conv2d(
class FullyConnectedLayer (line 70) | class FullyConnectedLayer(torch.nn.Module):
method __init__ (line 71) | def __init__(self,
method forward (line 90) | def forward(self, x):
method extra_repr (line 104) | def extra_repr(self):
class MappingNetwork (line 110) | class MappingNetwork(torch.nn.Module):
method __init__ (line 111) | def __init__(self,
method forward (line 136) | def forward(self, z, c, truncation_psi=1, truncation_cutoff=None, upda...
method extra_repr (line 164) | def extra_repr(self):
class SynthesisInput (line 170) | class SynthesisInput(torch.nn.Module):
method __init__ (line 171) | def __init__(self,
method forward (line 199) | def forward(self, w):
method extra_repr (line 246) | def extra_repr(self):
class SynthesisLayer (line 254) | class SynthesisLayer(torch.nn.Module):
method __init__ (line 255) | def __init__(self,
method forward (line 330) | def forward(self, x, w, noise_mode='random', force_fp32=False, update_...
method design_lowpass_filter (line 365) | def design_lowpass_filter(numtaps, cutoff, width, fs, radial=False):
method extra_repr (line 387) | def extra_repr(self):
class SynthesisNetwork (line 400) | class SynthesisNetwork(torch.nn.Module):
method __init__ (line 401) | def __init__(self,
method forward (line 465) | def forward(self, ws, return_feature=False, **layer_kwargs):
method extra_repr (line 486) | def extra_repr(self):
class Generator (line 496) | class Generator(torch.nn.Module):
method __init__ (line 497) | def __init__(self,
method forward (line 518) | def forward(self, z, c, truncation_psi=1, truncation_cutoff=None, upda...
function imresize (line 537) | def imresize(image, size):
FILE: training/training_loop.py
function setup_snapshot_image_grid (line 31) | def setup_snapshot_image_grid(training_set, random_seed=0):
function save_image_grid (line 70) | def save_image_grid(img, fname, drange, grid_size):
function training_loop (line 90) | def training_loop(
FILE: visualizer_drag.py
class Visualizer (line 29) | class Visualizer(imgui_window.ImguiWindow):
method __init__ (line 30) | def __init__(self, capture_dir=None):
method close (line 66) | def close(self):
method add_recent_pickle (line 72) | def add_recent_pickle(self, pkl, ignore_errors=False):
method load_pickle (line 75) | def load_pickle(self, pkl, ignore_errors=False):
method print_error (line 78) | def print_error(self, error):
method defer_rendering (line 84) | def defer_rendering(self, num_frames=1):
method clear_result (line 87) | def clear_result(self):
method set_async (line 90) | def set_async(self, is_async):
method _adjust_font_size (line 98) | def _adjust_font_size(self):
method check_update_mask (line 104) | def check_update_mask(self, **args):
method capture_image_frame (line 116) | def capture_image_frame(self):
method get_drag_info (line 125) | def get_drag_info(self):
method draw_frame (line 133) | def draw_frame(self):
class AsyncRenderer (line 253) | class AsyncRenderer:
method __init__ (line 254) | def __init__(self):
method close (line 265) | def close(self):
method is_async (line 275) | def is_async(self):
method set_async (line 278) | def set_async(self, is_async):
method set_args (line 281) | def set_args(self, **args):
method _set_args_async (line 298) | def _set_args_async(self, **args):
method _set_args_sync (line 310) | def _set_args_sync(self, **args):
method get_result (line 315) | def get_result(self):
method clear_result (line 324) | def clear_result(self):
method _process_fn (line 331) | def _process_fn(args_queue, result_queue):
function main (line 353) | def main(
FILE: visualizer_drag_gradio.py
function reverse_point_pairs (line 32) | def reverse_point_pairs(points):
function clear_state (line 39) | def clear_state(global_state, target=None):
function init_images (line 61) | def init_images(global_state):
function update_image_draw (line 102) | def update_image_draw(image, points, mask, show_mask, global_state=None):
function preprocess_mask_info (line 115) | def preprocess_mask_info(global_state, image):
function on_change_pretrained_dropdown (line 350) | def on_change_pretrained_dropdown(pretrained_value, global_state):
function on_click_reset_image (line 368) | def on_click_reset_image(global_state):
function on_change_update_image_seed (line 386) | def on_change_update_image_seed(seed, global_state):
function on_click_latent_space (line 404) | def on_click_latent_space(latent_space, global_state):
function on_change_lr (line 428) | def on_change_lr(lr, global_state):
function on_click_start (line 446) | def on_click_start(global_state, image):
function on_click_stop (line 672) | def on_click_stop(global_state):
function on_click_remove_point (line 695) | def on_click_remove_point(global_state):
function on_click_reset_mask (line 710) | def on_click_reset_mask(global_state):
function on_click_enable_draw (line 731) | def on_click_enable_draw(global_state, image):
function on_click_remove_draw (line 746) | def on_click_remove_draw(global_state, image):
function on_click_add_point (line 768) | def on_click_add_point(global_state, image: dict):
function on_click_image (line 789) | def on_click_image(global_state, evt: gr.SelectData):
function on_click_clear_points (line 829) | def on_click_clear_points(global_state):
function on_click_show_mask (line 849) | def on_click_show_mask(global_state, show_mask):
FILE: viz/capture_widget.py
class CaptureWidget (line 21) | class CaptureWidget:
method __init__ (line 22) | def __init__(self, viz):
method dump_png (line 30) | def dump_png(self, image):
method __call__ (line 53) | def __call__(self, show=True):
FILE: viz/drag_widget.py
class DragWidget (line 10) | class DragWidget:
method __init__ (line 11) | def __init__(self, viz):
method action (line 32) | def action(self, click, down, x, y):
method add_point (line 38) | def add_point(self, click, x, y):
method init_mask (line 52) | def init_mask(self, w, h):
method draw_mask (line 56) | def draw_mask(self, x, y):
method stop_drag (line 66) | def stop_drag(self):
method set_points (line 70) | def set_points(self, points):
method reset_point (line 73) | def reset_point(self):
method load_points (line 78) | def load_points(self, suffix):
method __call__ (line 91) | def __call__(self, show=True):
FILE: viz/latent_widget.py
class LatentWidget (line 18) | class LatentWidget:
method __init__ (line 19) | def __init__(self, viz):
method __call__ (line 31) | def __call__(self, show=True):
FILE: viz/pickle_widget.py
function _locate_results (line 22) | def _locate_results(pattern):
class PickleWidget (line 27) | class PickleWidget:
method __init__ (line 28) | def __init__(self, viz):
method add_recent (line 38) | def add_recent(self, pkl, ignore_errors=False):
method load (line 47) | def load(self, pkl, ignore_errors=False):
method __call__ (line 72) | def __call__(self, show=True):
method list_runs_and_pkls (line 136) | def list_runs_and_pkls(self, parents):
method resolve_pkl (line 151) | def resolve_pkl(self, pattern):
FILE: viz/renderer.py
class CapturedException (line 27) | class CapturedException(Exception):
method __init__ (line 28) | def __init__(self, msg=None):
class CaptureSuccess (line 41) | class CaptureSuccess(Exception):
method __init__ (line 42) | def __init__(self, out):
function add_watermark_np (line 48) | def add_watermark_np(input_image_array, watermark_text="AI Generated"):
class Renderer (line 70) | class Renderer:
method __init__ (line 71) | def __init__(self, disable_timing=False):
method render (line 85) | def render(self, **args):
method get_network (line 133) | def get_network(self, pkl, key, **tweak_kwargs):
method _get_pinned_buf (line 179) | def _get_pinned_buf(self, ref):
method to_device (line 187) | def to_device(self, buf):
method to_cpu (line 190) | def to_cpu(self, buf):
method _ignore_timing (line 193) | def _ignore_timing(self):
method _apply_cmap (line 196) | def _apply_cmap(self, x, name='viridis'):
method init_network (line 208) | def init_network(self, res,
method update_lr (line 265) | def update_lr(self, lr):
method _render_drag_impl (line 272) | def _render_drag_impl(self, res,
Condensed preview — 186 files, each showing path, character count, and a content snippet. Download the .json file or copy for the full structured content (1,642K chars).
[
{
"path": ".gitignore",
"chars": 2281,
"preview": "# Created by .ignore support plugin (hsz.mobi)\n### Python template\n# Byte-compiled / optimized / DLL files\n__pycache__/\n"
},
{
"path": "Dockerfile",
"chars": 689,
"preview": "FROM nvcr.io/nvidia/pytorch:23.05-py3\n\nENV PYTHONDONTWRITEBYTECODE 1\nENV PYTHONUNBUFFERED 1\n\nRUN apt-get update && apt-g"
},
{
"path": "LICENSE.txt",
"chars": 4383,
"preview": "Copyright (c) 2021, NVIDIA Corporation & affiliates. All rights reserved.\n\n\nNVIDIA Source Code License for StyleGAN3\n\n\n="
},
{
"path": "README.md",
"chars": 5884,
"preview": "<p align=\"center\">\n\n <h1 align=\"center\">Drag Your GAN: Interactive Point-based Manipulation on the Generative Image Man"
},
{
"path": "dnnlib/__init__.py",
"chars": 488,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "dnnlib/util.py",
"chars": 17132,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "environment.yml",
"chars": 550,
"preview": "name: stylegan3\nchannels:\n - pytorch\n - nvidia\ndependencies:\n - python >= 3.8\n - pip\n - numpy>=1.25\n - click>=8.0\n"
},
{
"path": "gen_images.py",
"chars": 6091,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "gradio_utils/__init__.py",
"chars": 346,
"preview": "from .utils import (ImageMask, draw_mask_on_image, draw_points_on_image,\n get_latest_points_pair, get"
},
{
"path": "gradio_utils/utils.py",
"chars": 4856,
"preview": "import gradio as gr\nimport numpy as np\nfrom PIL import Image, ImageDraw\n\n\nclass ImageMask(gr.components.Image):\n \"\"\"\n"
},
{
"path": "gui_utils/__init__.py",
"chars": 448,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "gui_utils/gl_utils.py",
"chars": 17409,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "gui_utils/glfw_window.py",
"chars": 7848,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "gui_utils/imgui_utils.py",
"chars": 7633,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "gui_utils/imgui_window.py",
"chars": 3994,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "gui_utils/text_utils.py",
"chars": 5637,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "legacy.py",
"chars": 16561,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "requirements.txt",
"chars": 182,
"preview": "torch>=2.0.0\nscipy>=1.11.1\nNinja==1.10.2\ngradio>=3.35.2\nimageio-ffmpeg>=0.4.3\nhuggingface_hub\nhf_transfer\npyopengl\nimgui"
},
{
"path": "stylegan_human/.gitignore",
"chars": 85,
"preview": ".DS_Store\n__pycache__\n*.pt\n*.pth\n*.pdparams\n*.pdiparams\n*.pdmodel\n*.pkl\n*.info\n*.yaml"
},
{
"path": "stylegan_human/PP_HumanSeg/deploy/infer.py",
"chars": 6331,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\n\n# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve"
},
{
"path": "stylegan_human/PP_HumanSeg/export_model/download_export_model.py",
"chars": 1849,
"preview": "# coding: utf8\n# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.\n#\n# Licensed under the Apache License, Ver"
},
{
"path": "stylegan_human/PP_HumanSeg/pretrained_model/download_pretrained_model.py",
"chars": 1743,
"preview": "# coding: utf8\n# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve.\n#\n# Licensed under the Apache License, Ver"
},
{
"path": "stylegan_human/README.md",
"chars": 15131,
"preview": "# StyleGAN-Human: A Data-Centric Odyssey of Human Generation\n<img src=\"./img/demo_V5_thumbnails-min.png\" width=\"96%\" he"
},
{
"path": "stylegan_human/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "stylegan_human/alignment.py",
"chars": 9314,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\n\nimport os\nimport argparse \nimport numpy as np\nimport torch\nfr"
},
{
"path": "stylegan_human/bg_white.py",
"chars": 1771,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\nimport os\nimport click\nimport cv2\nimport numpy as np\n\ndef bg_w"
},
{
"path": "stylegan_human/dnnlib/__init__.py",
"chars": 534,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\n# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserve"
},
{
"path": "stylegan_human/dnnlib/tflib/__init__.py",
"chars": 525,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\n# Copyright (c) 2019, NVIDIA Corporation. All rights reserved."
},
{
"path": "stylegan_human/dnnlib/tflib/autosummary.py",
"chars": 8096,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\n# Copyright (c) 2019, NVIDIA Corporation. All rights reserved."
},
{
"path": "stylegan_human/dnnlib/tflib/custom_ops.py",
"chars": 7281,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\n# Copyright (c) 2019, NVIDIA Corporation. All rights reserved."
},
{
"path": "stylegan_human/dnnlib/tflib/network.py",
"chars": 30040,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\n# Copyright (c) 2019, NVIDIA Corporation. All rights reserved."
},
{
"path": "stylegan_human/dnnlib/tflib/ops/__init__.py",
"chars": 293,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\n# Copyright (c) 2019, NVIDIA Corporation. All rights reserved."
},
{
"path": "stylegan_human/dnnlib/tflib/ops/fused_bias_act.cu",
"chars": 7851,
"preview": "// Copyright (c) SenseTime Research. All rights reserved.\n\n// Copyright (c) 2019, NVIDIA Corporation. All rights reserve"
},
{
"path": "stylegan_human/dnnlib/tflib/ops/fused_bias_act.py",
"chars": 8597,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\n# Copyright (c) 2019, NVIDIA Corporation. All rights reserved."
},
{
"path": "stylegan_human/dnnlib/tflib/ops/upfirdn_2d.cu",
"chars": 15201,
"preview": "// Copyright (c) SenseTime Research. All rights reserved.\n\n// Copyright (c) 2019, NVIDIA Corporation. All rights reserve"
},
{
"path": "stylegan_human/dnnlib/tflib/ops/upfirdn_2d.py",
"chars": 16646,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\n# Copyright (c) 2019, NVIDIA Corporation. All rights reserved."
},
{
"path": "stylegan_human/dnnlib/tflib/optimizer.py",
"chars": 18027,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\n# Copyright (c) 2019, NVIDIA Corporation. All rights reserved."
},
{
"path": "stylegan_human/dnnlib/tflib/tfutil.py",
"chars": 9733,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\n# Copyright (c) 2019, NVIDIA Corporation. All rights reserved."
},
{
"path": "stylegan_human/dnnlib/util.py",
"chars": 16760,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved"
},
{
"path": "stylegan_human/docs/Dataset.md",
"chars": 4498,
"preview": "# SHHQ Dataset\n<img src=\"../img/preview_samples1.png\" width=\"96%\" height=\"96%\">\n\n## Overview\nSHHQ is a dataset with high"
},
{
"path": "stylegan_human/edit/__init__.py",
"chars": 65,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\n# empty"
},
{
"path": "stylegan_human/edit/edit_config.py",
"chars": 596,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\nattr_dict = dict(\n interface_gan={ # strength\n 'uppe"
},
{
"path": "stylegan_human/edit/edit_helper.py",
"chars": 8147,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\nfrom legacy import save_obj, load_pkl\nimport torch\nfrom torch."
},
{
"path": "stylegan_human/edit.py",
"chars": 9579,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\nimport os\nimport sys\nimport torch\nimport numpy as np\nsys.path."
},
{
"path": "stylegan_human/environment.yml",
"chars": 484,
"preview": "name: stylehuman\nchannels:\n - pytorch\n - nvidia\ndependencies:\n - python == 3.8\n - pip\n - numpy>=1.20\n - click>=8.0"
},
{
"path": "stylegan_human/generate.py",
"chars": 4709,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\n# Copyright (c) 2019, NVIDIA Corporation. All rights reserved."
},
{
"path": "stylegan_human/insetgan.py",
"chars": 20022,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\n\nimport torch\nimport torch.nn.functional as F\nfrom tqdm import"
},
{
"path": "stylegan_human/interpolation.py",
"chars": 5616,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\n## interpolate between two z code\n## score all middle latent c"
},
{
"path": "stylegan_human/latent_direction/ss_statics/bottom_length_statis/3/statis.csv",
"chars": 10793,
"preview": ",channel,strength\n401,401,0.0051189866\n79,79,0.004417926\n499,499,0.0042351373\n272,272,0.0033855115\n2,2,0.003143758\n267,2"
},
{
"path": "stylegan_human/latent_direction/ss_statics/bottom_length_statis/4/statis.csv",
"chars": 10605,
"preview": ",channel,strength\n371,371,0.010705759\n130,130,0.007931795\n66,66,0.0069621187\n411,411,0.0065370337\n241,241,0.0061685536\n1"
},
{
"path": "stylegan_human/latent_direction/ss_statics/bottom_length_statis/5/statis.csv",
"chars": 10755,
"preview": ",channel,strength\n242,242,0.01746412\n134,134,0.011444086\n71,71,0.01060778\n395,395,0.0062382165\n363,363,0.0058679837\n175,"
},
{
"path": "stylegan_human/latent_direction/ss_statics/upper_length_statis/5/statis.csv",
"chars": 10846,
"preview": ",channel,strength\n423,423,0.004408924\n341,341,0.0032079767\n379,379,0.0028457695\n184,184,0.0028266357\n368,368,0.002722497"
},
{
"path": "stylegan_human/legacy.py",
"chars": 8173,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\n# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserve"
},
{
"path": "stylegan_human/openpose/model/.gitkeep",
"chars": 0,
"preview": ""
},
{
"path": "stylegan_human/openpose/src/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "stylegan_human/openpose/src/body.py",
"chars": 10990,
"preview": "import cv2\nimport numpy as np\nimport math\nimport time\nfrom scipy.ndimage.filters import gaussian_filter\nimport matplotli"
},
{
"path": "stylegan_human/openpose/src/model.py",
"chars": 8745,
"preview": "import torch\nfrom collections import OrderedDict\n\nimport torch\nimport torch.nn as nn\n\ndef make_layers(block, no_relu_lay"
},
{
"path": "stylegan_human/openpose/src/util.py",
"chars": 3750,
"preview": "import numpy as np\nimport math\nimport cv2\nimport matplotlib\nfrom matplotlib.backends.backend_agg import FigureCanvasAgg "
},
{
"path": "stylegan_human/pti/pti_configs/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "stylegan_human/pti/pti_configs/global_config.py",
"chars": 230,
"preview": "## Device\ncuda_visible_devices = '0'\ndevice = 'cuda:0'\n\n## Logs\ntraining_step = 1\nimage_rec_result_log_snapshot = 100\npi"
},
{
"path": "stylegan_human/pti/pti_configs/hyperparameters.py",
"chars": 554,
"preview": "## Architechture\nlpips_type = 'alex'\nfirst_inv_type = 'w+'#'w+'\noptim_type = 'adam'\n\n## Locality regularization\nlatent_b"
},
{
"path": "stylegan_human/pti/pti_configs/paths_config.py",
"chars": 726,
"preview": "import os\n\n## Pretrained models paths\ne4e = './pti/e4e_w+.pt'\nstylegan2_ada_shhq = './pretrained_models/stylegan_human_v"
},
{
"path": "stylegan_human/pti/pti_models/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "stylegan_human/pti/pti_models/e4e/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "stylegan_human/pti/pti_models/e4e/encoders/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "stylegan_human/pti/pti_models/e4e/encoders/helpers.py",
"chars": 4882,
"preview": "from collections import namedtuple\nimport torch\nimport torch.nn.functional as F\nfrom torch.nn import Conv2d, BatchNorm2d"
},
{
"path": "stylegan_human/pti/pti_models/e4e/encoders/model_irse.py",
"chars": 3257,
"preview": "from torch.nn import Linear, Conv2d, BatchNorm1d, BatchNorm2d, PReLU, Dropout, Sequential, Module\nfrom encoder4editing.m"
},
{
"path": "stylegan_human/pti/pti_models/e4e/encoders/psp_encoders.py",
"chars": 7204,
"preview": "from enum import Enum\nimport math\nimport numpy as np\nimport torch\nfrom torch import nn\nfrom torch.nn import Conv2d, Batc"
},
{
"path": "stylegan_human/pti/pti_models/e4e/latent_codes_pool.py",
"chars": 2349,
"preview": "import random\nimport torch\n\n\nclass LatentCodesPool:\n \"\"\"This class implements latent codes buffer that stores previou"
},
{
"path": "stylegan_human/pti/pti_models/e4e/psp.py",
"chars": 4142,
"preview": "import matplotlib\nfrom pti.pti_configs import paths_config\nmatplotlib.use('Agg')\nimport torch\nfrom torch import nn\nfrom "
},
{
"path": "stylegan_human/pti/pti_models/e4e/stylegan2/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "stylegan_human/pti/pti_models/e4e/stylegan2/model.py",
"chars": 18968,
"preview": "import math\nimport random\nimport torch\nfrom torch import nn\nfrom torch.nn import functional as F\n\nfrom .op.fused_act imp"
},
{
"path": "stylegan_human/pti/pti_models/e4e/stylegan2/op/__init__.py",
"chars": 89,
"preview": "from .fused_act import FusedLeakyReLU, fused_leaky_relu\nfrom .upfirdn2d import upfirdn2d\n"
},
{
"path": "stylegan_human/pti/pti_models/e4e/stylegan2/op/fused_act.py",
"chars": 840,
"preview": "import os\n\nimport torch\nfrom torch import nn\nfrom torch.nn import functional as F\nfrom torch.autograd import Function\n\n\n"
},
{
"path": "stylegan_human/pti/pti_models/e4e/stylegan2/op/fused_bias_act.cpp",
"chars": 826,
"preview": "#include <torch/extension.h>\n\n\ntorch::Tensor fused_bias_act_op(const torch::Tensor& input, const torch::Tensor& bias, co"
},
{
"path": "stylegan_human/pti/pti_models/e4e/stylegan2/op/fused_bias_act_kernel.cu",
"chars": 2777,
"preview": "// Copyright (c) 2019, NVIDIA Corporation. All rights reserved.\n//\n// This work is made available under the Nvidia Sourc"
},
{
"path": "stylegan_human/pti/pti_models/e4e/stylegan2/op/upfirdn2d.cpp",
"chars": 966,
"preview": "#include <torch/extension.h>\n\n\ntorch::Tensor upfirdn2d_op(const torch::Tensor& input, const torch::Tensor& kernel,\n "
},
{
"path": "stylegan_human/pti/pti_models/e4e/stylegan2/op/upfirdn2d.py",
"chars": 1649,
"preview": "import os\n\nimport torch\nfrom torch.nn import functional as F\n\n\nmodule_path = os.path.dirname(__file__)\n\n\n\ndef upfirdn2d("
},
{
"path": "stylegan_human/pti/pti_models/e4e/stylegan2/op/upfirdn2d_kernel.cu",
"chars": 8953,
"preview": "// Copyright (c) 2019, NVIDIA Corporation. All rights reserved.\n//\n// This work is made available under the Nvidia Sourc"
},
{
"path": "stylegan_human/pti/training/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "stylegan_human/pti/training/coaches/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "stylegan_human/pti/training/coaches/base_coach.py",
"chars": 5755,
"preview": "import abc\nimport os\nimport pickle\nfrom argparse import Namespace\nimport wandb\nimport os.path\nfrom .localitly_regulizer "
},
{
"path": "stylegan_human/pti/training/coaches/localitly_regulizer.py",
"chars": 2946,
"preview": "import torch\nimport numpy as np\nimport wandb\nfrom pti.pti_configs import hyperparameters, global_config\nl2_criterion = t"
},
{
"path": "stylegan_human/pti/training/coaches/multi_id_coach.py",
"chars": 2902,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\nimport os\n\nimport torch\nfrom tqdm import tqdm\n\nfrom pti.pti_co"
},
{
"path": "stylegan_human/pti/training/coaches/single_id_coach.py",
"chars": 3471,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\nimport os\nimport torch\nfrom tqdm import tqdm\nfrom pti.pti_conf"
},
{
"path": "stylegan_human/pti/training/projectors/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "stylegan_human/pti/training/projectors/w_plus_projector.py",
"chars": 5939,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors retain all inte"
},
{
"path": "stylegan_human/pti/training/projectors/w_projector.py",
"chars": 5977,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors retain all inte"
},
{
"path": "stylegan_human/run_pti.py",
"chars": 1873,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\nfrom random import choice\nfrom string import ascii_uppercase\nf"
},
{
"path": "stylegan_human/style_mixing.py",
"chars": 4579,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\n# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved"
},
{
"path": "stylegan_human/stylemixing_video.py",
"chars": 7444,
"preview": "\n# Copyright (c) SenseTime Research. All rights reserved.\n\n\"\"\"Here we demo style-mixing results using StyleGAN2 pretrain"
},
{
"path": "stylegan_human/torch_utils/__init__.py",
"chars": 494,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\n# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserve"
},
{
"path": "stylegan_human/torch_utils/custom_ops.py",
"chars": 10603,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\n# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved"
},
{
"path": "stylegan_human/torch_utils/misc.py",
"chars": 11050,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\n# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserve"
},
{
"path": "stylegan_human/torch_utils/models.py",
"chars": 22534,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\n# https://github.com/rosinality/stylegan2-pytorch/blob/master/"
},
{
"path": "stylegan_human/torch_utils/models_face.py",
"chars": 22703,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\nimport math\nimport random\nimport functools\nimport operator\n\nim"
},
{
"path": "stylegan_human/torch_utils/op_edit/__init__.py",
"chars": 147,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\nfrom .fused_act import FusedLeakyReLU, fused_leaky_relu\nfrom ."
},
{
"path": "stylegan_human/torch_utils/op_edit/fused_act.py",
"chars": 2752,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\nimport os\n\nimport torch\nfrom torch import nn\nfrom torch.nn imp"
},
{
"path": "stylegan_human/torch_utils/op_edit/fused_bias_act.cpp",
"chars": 885,
"preview": "// Copyright (c) SenseTime Research. All rights reserved.\n\n#include <torch/extension.h>\n\n\ntorch::Tensor fused_bias_act_o"
},
{
"path": "stylegan_human/torch_utils/op_edit/fused_bias_act_kernel.cu",
"chars": 2936,
"preview": "// Copyright (c) SenseTime Research. All rights reserved.\r\n\r\n// Copyright (c) 2019, NVIDIA Corporation. All rights reser"
},
{
"path": "stylegan_human/torch_utils/op_edit/upfirdn2d.cpp",
"chars": 1025,
"preview": "// Copyright (c) SenseTime Research. All rights reserved.\n\n#include <torch/extension.h>\n\n\ntorch::Tensor upfirdn2d_op(con"
},
{
"path": "stylegan_human/torch_utils/op_edit/upfirdn2d.py",
"chars": 5732,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\nimport os\n\nimport torch\nfrom torch.nn import functional as F\nf"
},
{
"path": "stylegan_human/torch_utils/op_edit/upfirdn2d_kernel.cu",
"chars": 12140,
"preview": "// Copyright (c) SenseTime Research. All rights reserved.\r\n\r\n// Copyright (c) 2019, NVIDIA Corporation. All rights reser"
},
{
"path": "stylegan_human/torch_utils/ops/__init__.py",
"chars": 65,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\n#empty"
},
{
"path": "stylegan_human/torch_utils/ops/bias_act.cpp",
"chars": 4435,
"preview": "// Copyright (c) SenseTime Research. All rights reserved.\n\n// Copyright (c) 2021, NVIDIA CORPORATION. All rights reserv"
},
{
"path": "stylegan_human/torch_utils/ops/bias_act.cu",
"chars": 6206,
"preview": "// Copyright (c) SenseTime Research. All rights reserved.\n\n// Copyright (c) 2021, NVIDIA CORPORATION. All rights reserv"
},
{
"path": "stylegan_human/torch_utils/ops/bias_act.h",
"chars": 1339,
"preview": "// Copyright (c) SenseTime Research. All rights reserved.\n\n// Copyright (c) 2021, NVIDIA CORPORATION. All rights reserv"
},
{
"path": "stylegan_human/torch_utils/ops/bias_act.py",
"chars": 10105,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\n# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved"
},
{
"path": "stylegan_human/torch_utils/ops/conv2d_gradfix.py",
"chars": 7735,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\n# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved"
},
{
"path": "stylegan_human/torch_utils/ops/conv2d_resample.py",
"chars": 7649,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\n# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved"
},
{
"path": "stylegan_human/torch_utils/ops/filtered_lrelu.cpp",
"chars": 15596,
"preview": "// Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n//\n// NVIDIA CORPORATION and its licensors"
},
{
"path": "stylegan_human/torch_utils/ops/filtered_lrelu.cu",
"chars": 67306,
"preview": "// Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n//\n// NVIDIA CORPORATION and its licensors"
},
{
"path": "stylegan_human/torch_utils/ops/filtered_lrelu.h",
"chars": 4447,
"preview": "// Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n//\n// NVIDIA CORPORATION and its licensors"
},
{
"path": "stylegan_human/torch_utils/ops/filtered_lrelu.py",
"chars": 13520,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "stylegan_human/torch_utils/ops/filtered_lrelu_ns.cu",
"chars": 1637,
"preview": "// Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n//\n// NVIDIA CORPORATION and its licensors"
},
{
"path": "stylegan_human/torch_utils/ops/filtered_lrelu_rd.cu",
"chars": 1606,
"preview": "// Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n//\n// NVIDIA CORPORATION and its licensors"
},
{
"path": "stylegan_human/torch_utils/ops/filtered_lrelu_wr.cu",
"chars": 1607,
"preview": "// Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n//\n// NVIDIA CORPORATION and its licensors"
},
{
"path": "stylegan_human/torch_utils/ops/fma.py",
"chars": 2092,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\n# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved"
},
{
"path": "stylegan_human/torch_utils/ops/grid_sample_gradfix.py",
"chars": 3357,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\n# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved"
},
{
"path": "stylegan_human/torch_utils/ops/upfirdn2d.cpp",
"chars": 4618,
"preview": "// Copyright (c) SenseTime Research. All rights reserved.\n\n// Copyright (c) 2021, NVIDIA CORPORATION. All rights reserv"
},
{
"path": "stylegan_human/torch_utils/ops/upfirdn2d.cu",
"chars": 21109,
"preview": "// Copyright (c) SenseTime Research. All rights reserved.\n\n// Copyright (c) 2021, NVIDIA CORPORATION. All rights reserv"
},
{
"path": "stylegan_human/torch_utils/ops/upfirdn2d.h",
"chars": 1895,
"preview": "// Copyright (c) SenseTime Research. All rights reserved.\n\n// Copyright (c) 2021, NVIDIA CORPORATION. All rights reserv"
},
{
"path": "stylegan_human/torch_utils/ops/upfirdn2d.py",
"chars": 16345,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\n# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved"
},
{
"path": "stylegan_human/torch_utils/persistence.py",
"chars": 9766,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\n# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserve"
},
{
"path": "stylegan_human/torch_utils/training_stats.py",
"chars": 10765,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\n# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved"
},
{
"path": "stylegan_human/training/__init__.py",
"chars": 448,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "stylegan_human/training/augment.py",
"chars": 26617,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "stylegan_human/training/dataset.py",
"chars": 8642,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "stylegan_human/training/loss.py",
"chars": 7999,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "stylegan_human/training/networks_stylegan2.py",
"chars": 41296,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "stylegan_human/training/networks_stylegan3.py",
"chars": 26923,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "stylegan_human/training/training_loop.py",
"chars": 21879,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "stylegan_human/training_scripts/sg2/train.py",
"chars": 24807,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\n# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved"
},
{
"path": "stylegan_human/training_scripts/sg2/training/dataset.py",
"chars": 9230,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\n# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserve"
},
{
"path": "stylegan_human/training_scripts/sg2/training/networks.py",
"chars": 40796,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\n# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserve"
},
{
"path": "stylegan_human/training_scripts/sg3/train.py",
"chars": 16259,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\n# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All ri"
},
{
"path": "stylegan_human/training_scripts/sg3/training/dataset.py",
"chars": 9382,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\n# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved"
},
{
"path": "stylegan_human/training_scripts/sg3/training/networks_stylegan2.py",
"chars": 42758,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\n# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All ri"
},
{
"path": "stylegan_human/training_scripts/sg3/training/networks_stylegan3.py",
"chars": 27144,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\n# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All ri"
},
{
"path": "stylegan_human/utils/ImagesDataset.py",
"chars": 707,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\n\nimport os\nfrom torch.utils.data import Dataset\nfrom PIL impor"
},
{
"path": "stylegan_human/utils/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "stylegan_human/utils/data_utils.py",
"chars": 953,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\n\nimport os\n\nfrom PIL import Image\n\nIMG_EXTENSIONS = [\n '.jp"
},
{
"path": "stylegan_human/utils/face_alignment.py",
"chars": 9914,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\nimport numpy as np\nimport PIL\nimport PIL.Image\nimport scipy\nim"
},
{
"path": "stylegan_human/utils/log_utils.py",
"chars": 2519,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\n\nimport numpy as np\nfrom PIL import Image\nimport wandb\nfrom pt"
},
{
"path": "stylegan_human/utils/models_utils.py",
"chars": 741,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\n\nimport pickle\nimport functools\nimport torch\nfrom pti.pti_conf"
},
{
"path": "stylegan_human/utils/util.py",
"chars": 1998,
"preview": "# Copyright (c) SenseTime Research. All rights reserved.\n\nimport torch\nimport cv2\nfrom torchvision import transforms\nimp"
},
{
"path": "torch_utils/__init__.py",
"chars": 448,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "torch_utils/custom_ops.py",
"chars": 6646,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "torch_utils/misc.py",
"chars": 11106,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "torch_utils/ops/__init__.py",
"chars": 448,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "torch_utils/ops/bias_act.cpp",
"chars": 4389,
"preview": "// Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n//\n// NVIDIA CORPORATION and its licensors"
},
{
"path": "torch_utils/ops/bias_act.cu",
"chars": 6160,
"preview": "// Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n//\n// NVIDIA CORPORATION and its licensors"
},
{
"path": "torch_utils/ops/bias_act.h",
"chars": 1293,
"preview": "// Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n//\n// NVIDIA CORPORATION and its licensors"
},
{
"path": "torch_utils/ops/bias_act.py",
"chars": 9845,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "torch_utils/ops/conv2d_gradfix.py",
"chars": 9465,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "torch_utils/ops/conv2d_resample.py",
"chars": 6765,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "torch_utils/ops/filtered_lrelu.cpp",
"chars": 15597,
"preview": "// Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n//\n// NVIDIA CORPORATION and its licensors"
},
{
"path": "torch_utils/ops/filtered_lrelu.cu",
"chars": 67307,
"preview": "// Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n//\n// NVIDIA CORPORATION and its licensors"
},
{
"path": "torch_utils/ops/filtered_lrelu.h",
"chars": 4448,
"preview": "// Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n//\n// NVIDIA CORPORATION and its licensors"
},
{
"path": "torch_utils/ops/filtered_lrelu.py",
"chars": 12916,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "torch_utils/ops/filtered_lrelu_ns.cu",
"chars": 1638,
"preview": "// Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n//\n// NVIDIA CORPORATION and its licensors"
},
{
"path": "torch_utils/ops/filtered_lrelu_rd.cu",
"chars": 1607,
"preview": "// Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n//\n// NVIDIA CORPORATION and its licensors"
},
{
"path": "torch_utils/ops/filtered_lrelu_wr.cu",
"chars": 1608,
"preview": "// Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n//\n// NVIDIA CORPORATION and its licensors"
},
{
"path": "torch_utils/ops/fma.py",
"chars": 2047,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "torch_utils/ops/grid_sample_gradfix.py",
"chars": 3020,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "torch_utils/ops/upfirdn2d.cpp",
"chars": 5027,
"preview": "// Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n//\n// NVIDIA CORPORATION and its licensors"
},
{
"path": "torch_utils/ops/upfirdn2d.cu",
"chars": 23137,
"preview": "// Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n//\n// NVIDIA CORPORATION and its licensors"
},
{
"path": "torch_utils/ops/upfirdn2d.h",
"chars": 1849,
"preview": "// Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n//\n// NVIDIA CORPORATION and its licensors"
},
{
"path": "torch_utils/ops/upfirdn2d.py",
"chars": 16424,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "torch_utils/persistence.py",
"chars": 9752,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "torch_utils/training_stats.py",
"chars": 10720,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "training/__init__.py",
"chars": 448,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "training/augment.py",
"chars": 26617,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "training/dataset.py",
"chars": 8642,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "training/loss.py",
"chars": 7999,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "training/networks_stylegan2.py",
"chars": 41518,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "training/networks_stylegan3.py",
"chars": 27262,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "training/training_loop.py",
"chars": 21879,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "visualizer_drag.py",
"chars": 17150,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "visualizer_drag_gradio.py",
"chars": 32392,
"preview": "import os\nimport os.path as osp\nfrom argparse import ArgumentParser\nfrom functools import partial\n\nimport gradio as gr\ni"
},
{
"path": "viz/__init__.py",
"chars": 448,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "viz/capture_widget.py",
"chars": 3824,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "viz/drag_widget.py",
"chars": 6098,
"preview": "import os\nimport torch\nimport numpy as np\nimport imgui\nimport dnnlib\nfrom gui_utils import imgui_utils\n\n#---------------"
},
{
"path": "viz/latent_widget.py",
"chars": 3964,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "viz/pickle_widget.py",
"chars": 6898,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
},
{
"path": "viz/renderer.py",
"chars": 15576,
"preview": "# Copyright (c) 2021, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# NVIDIA CORPORATION and its licensors re"
}
]
About this extraction
This page contains the full source code of the XingangPan/DragGAN GitHub repository, extracted and formatted as plain text for AI agents and large language models (LLMs). The extraction includes 186 files (1.5 MB), approximately 427.1k tokens, and a symbol index with 1360 extracted functions, classes, methods, constants, and types. Use this with OpenClaw, Claude, ChatGPT, Cursor, Windsurf, or any other AI tool that accepts text input. You can copy the full output to your clipboard or download it as a .txt file.
Extracted by GitExtract — free GitHub repo to text converter for AI. Built by Nikandr Surkov.