Repository: JCBrouwer/maua-stylegan2
Branch: master
Commit: 7f9282141053
Files: 86
Total size: 560.8 KB

Directory structure:
gitextract_lesmzkyd/

├── .gitignore
├── LICENSE/
│   ├── LICENSE-AUDIOREACTIVE
│   ├── LICENSE-AUTOENCODER
│   ├── LICENSE-CONTRASTIVE-LEARNER
│   ├── LICENSE-FID
│   ├── LICENSE-LPIPS
│   ├── LICENSE-LUCIDRAINS
│   ├── LICENSE-NVIDIA
│   ├── LICENSE-ROSINALITY
│   └── LICENSE-VGG
├── README.md
├── accelerate/
│   ├── accelerate_inception.py
│   ├── accelerate_logcosh.py
│   └── accelerate_segnet.py
├── audioreactive/
│   ├── __init__.py
│   ├── bend.py
│   ├── examples/
│   │   ├── __init__.py
│   │   ├── default.py
│   │   ├── kelp.py
│   │   ├── tauceti.py
│   │   └── temper.py
│   ├── latent.py
│   ├── signal.py
│   └── util.py
├── augment.py
├── contrastive_learner.py
├── convert_weight.py
├── dataset.py
├── distributed.py
├── generate.py
├── generate_audiovisual.py
├── generate_video.py
├── gpu_profile.py
├── gpumon.py
├── lightning.py
├── lookahead_minimax.py
├── lucidrains.py
├── models/
│   ├── autoencoder.py
│   ├── stylegan1.py
│   └── stylegan2.py
├── op/
│   ├── __init__.py
│   ├── fused_act.py
│   ├── fused_bias_act.cpp
│   ├── fused_bias_act_kernel.cu
│   ├── upfirdn2d.cpp
│   ├── upfirdn2d.py
│   └── upfirdn2d_kernel.cu
├── prepare_data.py
├── prepare_vae_codes.py
├── projector.py
├── render.py
├── requirements.txt
├── select_latents.py
├── train.py
├── train_profile.py
├── validation/
│   ├── __init__.py
│   ├── calc_fid.py
│   ├── calc_inception.py
│   ├── calc_ppl.py
│   ├── inception.py
│   ├── lpips/
│   │   ├── __init__.py
│   │   ├── base_model.py
│   │   ├── dist_model.py
│   │   ├── networks_basic.py
│   │   ├── pretrained_networks.py
│   │   ├── util.py
│   │   └── weights/
│   │       ├── v0.0/
│   │       │   ├── alex.pth
│   │       │   ├── squeeze.pth
│   │       │   └── vgg.pth
│   │       └── v0.1/
│   │           ├── alex.pth
│   │           ├── squeeze.pth
│   │           └── vgg.pth
│   ├── metrics.py
│   └── spectral_norm.py
└── workspace/
    ├── naamloos_average_pitch.npy
    ├── naamloos_bass_sum.npy
    ├── naamloos_drop_latents.npy
    ├── naamloos_drop_latents_1.npy
    ├── naamloos_high_average_pitch.npy
    ├── naamloos_high_pitches_mean.npy
    ├── naamloos_intro_latents.npy
    ├── naamloos_metadata.json
    ├── naamloos_onsets.npy
    ├── naamloos_params.json
    ├── naamloos_pitches_mean.npy
    └── naamloos_rms.npy

================================================
FILE CONTENTS
================================================

================================================
FILE: .gitignore
================================================
pretrained_models/
wandb
wandb/
*.lmdb/
*.pkl
checkpoints/
maua-stylegan/
.vscode
output/
workspace/*
!workspace
output/*
!output

# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
*$py.class

# C extensions
*.so

# Distribution / packaging
.Python
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
wheels/
pip-wheel-metadata/
share/python-wheels/
*.egg-info/
.installed.cfg
*.egg
MANIFEST

# PyInstaller
#  Usually these files are written by a python script from a template
#  before PyInstaller builds the exe, so as to inject date/other infos into it.
*.manifest
*.spec

# Installer logs
pip-log.txt
pip-delete-this-directory.txt

# Unit test / coverage reports
htmlcov/
.tox/
.nox/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*.cover
*.py,cover
.hypothesis/
.pytest_cache/

# Translations
*.mo
*.pot

# Django stuff:
*.log
local_settings.py
db.sqlite3
db.sqlite3-journal

# Flask stuff:
instance/
.webassets-cache

# Scrapy stuff:
.scrapy

# Sphinx documentation
docs/_build/

# PyBuilder
target/

# Jupyter Notebook
.ipynb_checkpoints

# IPython
profile_default/
ipython_config.py

# pyenv
.python-version

# pipenv
#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
#   However, in case of collaboration, if having platform-specific dependencies or dependencies
#   having no cross-platform support, pipenv may install dependencies that don't work, or not
#   install all needed dependencies.
#Pipfile.lock

# PEP 582; used by e.g. github.com/David-OConnor/pyflow
__pypackages__/

# Celery stuff
celerybeat-schedule
celerybeat.pid

# SageMath parsed files
*.sage.py

# Environments
.env
.venv
env/
venv/
ENV/
env.bak/
venv.bak/

# Spyder project settings
.spyderproject
.spyproject

# Rope project settings
.ropeproject

# mkdocs documentation
/site

# mypy
.mypy_cache/
.dmypy.json
dmypy.json

# Pyre type checker
.pyre/



================================================
FILE: LICENSE/LICENSE-AUDIOREACTIVE
================================================
Code for Audio-reactive Latent Interpolations with StyleGAN
Including the folder audioreactive/, generate_audiovisual.py, generate_video.py, select_latents.py, and render.py

Copyright (C) 2020 Hans Brouwer

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

TERMS AND CONDITIONS
0. Definitions.

“This License” refers to version 3 of the GNU General Public License.

“Copyright” also means copyright-like laws that apply to other kinds of works, such as semiconductor masks.

“The Program” refers to any copyrightable work licensed under this License. Each licensee is addressed as “you”. “Licensees” and “recipients” may be individuals or organizations.

To “modify” a work means to copy from or adapt all or part of the work in a fashion requiring copyright permission, other than the making of an exact copy. The resulting work is called a “modified version” of the earlier work or a work “based on” the earlier work.

A “covered work” means either the unmodified Program or a work based on the Program.

To “propagate” a work means to do anything with it that, without permission, would make you directly or secondarily liable for infringement under applicable copyright law, except executing it on a computer or modifying a private copy. Propagation includes copying, distribution (with or without modification), making available to the public, and in some countries other activities as well.

To “convey” a work means any kind of propagation that enables other parties to make or receive copies. Mere interaction with a user through a computer network, with no transfer of a copy, is not conveying.

An interactive user interface displays “Appropriate Legal Notices” to the extent that it includes a convenient and prominently visible feature that (1) displays an appropriate copyright notice, and (2) tells the user that there is no warranty for the work (except to the extent that warranties are provided), that licensees may convey the work under this License, and how to view a copy of this License. If the interface presents a list of user commands or options, such as a menu, a prominent item in the list meets this criterion.
1. Source Code.

The “source code” for a work means the preferred form of the work for making modifications to it. “Object code” means any non-source form of a work.

A “Standard Interface” means an interface that either is an official standard defined by a recognized standards body, or, in the case of interfaces specified for a particular programming language, one that is widely used among developers working in that language.

The “System Libraries” of an executable work include anything, other than the work as a whole, that (a) is included in the normal form of packaging a Major Component, but which is not part of that Major Component, and (b) serves only to enable use of the work with that Major Component, or to implement a Standard Interface for which an implementation is available to the public in source code form. A “Major Component”, in this context, means a major essential component (kernel, window system, and so on) of the specific operating system (if any) on which the executable work runs, or a compiler used to produce the work, or an object code interpreter used to run it.

The “Corresponding Source” for a work in object code form means all the source code needed to generate, install, and (for an executable work) run the object code and to modify the work, including scripts to control those activities. However, it does not include the work's System Libraries, or general-purpose tools or generally available free programs which are used unmodified in performing those activities but which are not part of the work. For example, Corresponding Source includes interface definition files associated with source files for the work, and the source code for shared libraries and dynamically linked subprograms that the work is specifically designed to require, such as by intimate data communication or control flow between those subprograms and other parts of the work.

The Corresponding Source need not include anything that users can regenerate automatically from other parts of the Corresponding Source.

The Corresponding Source for a work in source code form is that same work.
2. Basic Permissions.

All rights granted under this License are granted for the term of copyright on the Program, and are irrevocable provided the stated conditions are met. This License explicitly affirms your unlimited permission to run the unmodified Program. The output from running a covered work is covered by this License only if the output, given its content, constitutes a covered work. This License acknowledges your rights of fair use or other equivalent, as provided by copyright law.

You may make, run and propagate covered works that you do not convey, without conditions so long as your license otherwise remains in force. You may convey covered works to others for the sole purpose of having them make modifications exclusively for you, or provide you with facilities for running those works, provided that you comply with the terms of this License in conveying all material for which you do not control copyright. Those thus making or running the covered works for you must do so exclusively on your behalf, under your direction and control, on terms that prohibit them from making any copies of your copyrighted material outside their relationship with you.

Conveying under any other circumstances is permitted solely under the conditions stated below. Sublicensing is not allowed; section 10 makes it unnecessary.
3. Protecting Users' Legal Rights From Anti-Circumvention Law.

No covered work shall be deemed part of an effective technological measure under any applicable law fulfilling obligations under article 11 of the WIPO copyright treaty adopted on 20 December 1996, or similar laws prohibiting or restricting circumvention of such measures.

When you convey a covered work, you waive any legal power to forbid circumvention of technological measures to the extent such circumvention is effected by exercising rights under this License with respect to the covered work, and you disclaim any intention to limit operation or modification of the work as a means of enforcing, against the work's users, your or third parties' legal rights to forbid circumvention of technological measures.
4. Conveying Verbatim Copies.

You may convey verbatim copies of the Program's source code as you receive it, in any medium, provided that you conspicuously and appropriately publish on each copy an appropriate copyright notice; keep intact all notices stating that this License and any non-permissive terms added in accord with section 7 apply to the code; keep intact all notices of the absence of any warranty; and give all recipients a copy of this License along with the Program.

You may charge any price or no price for each copy that you convey, and you may offer support or warranty protection for a fee.
5. Conveying Modified Source Versions.

You may convey a work based on the Program, or the modifications to produce it from the Program, in the form of source code under the terms of section 4, provided that you also meet all of these conditions:

    a) The work must carry prominent notices stating that you modified it, and giving a relevant date.
    b) The work must carry prominent notices stating that it is released under this License and any conditions added under section 7. This requirement modifies the requirement in section 4 to “keep intact all notices”.
    c) You must license the entire work, as a whole, under this License to anyone who comes into possession of a copy. This License will therefore apply, along with any applicable section 7 additional terms, to the whole of the work, and all its parts, regardless of how they are packaged. This License gives no permission to license the work in any other way, but it does not invalidate such permission if you have separately received it.
    d) If the work has interactive user interfaces, each must display Appropriate Legal Notices; however, if the Program has interactive interfaces that do not display Appropriate Legal Notices, your work need not make them do so.

A compilation of a covered work with other separate and independent works, which are not by their nature extensions of the covered work, and which are not combined with it such as to form a larger program, in or on a volume of a storage or distribution medium, is called an “aggregate” if the compilation and its resulting copyright are not used to limit the access or legal rights of the compilation's users beyond what the individual works permit. Inclusion of a covered work in an aggregate does not cause this License to apply to the other parts of the aggregate.
6. Conveying Non-Source Forms.

You may convey a covered work in object code form under the terms of sections 4 and 5, provided that you also convey the machine-readable Corresponding Source under the terms of this License, in one of these ways:

    a) Convey the object code in, or embodied in, a physical product (including a physical distribution medium), accompanied by the Corresponding Source fixed on a durable physical medium customarily used for software interchange.
    b) Convey the object code in, or embodied in, a physical product (including a physical distribution medium), accompanied by a written offer, valid for at least three years and valid for as long as you offer spare parts or customer support for that product model, to give anyone who possesses the object code either (1) a copy of the Corresponding Source for all the software in the product that is covered by this License, on a durable physical medium customarily used for software interchange, for a price no more than your reasonable cost of physically performing this conveying of source, or (2) access to copy the Corresponding Source from a network server at no charge.
    c) Convey individual copies of the object code with a copy of the written offer to provide the Corresponding Source. This alternative is allowed only occasionally and noncommercially, and only if you received the object code with such an offer, in accord with subsection 6b.
    d) Convey the object code by offering access from a designated place (gratis or for a charge), and offer equivalent access to the Corresponding Source in the same way through the same place at no further charge. You need not require recipients to copy the Corresponding Source along with the object code. If the place to copy the object code is a network server, the Corresponding Source may be on a different server (operated by you or a third party) that supports equivalent copying facilities, provided you maintain clear directions next to the object code saying where to find the Corresponding Source. Regardless of what server hosts the Corresponding Source, you remain obligated to ensure that it is available for as long as needed to satisfy these requirements.
    e) Convey the object code using peer-to-peer transmission, provided you inform other peers where the object code and Corresponding Source of the work are being offered to the general public at no charge under subsection 6d.

A separable portion of the object code, whose source code is excluded from the Corresponding Source as a System Library, need not be included in conveying the object code work.

A “User Product” is either (1) a “consumer product”, which means any tangible personal property which is normally used for personal, family, or household purposes, or (2) anything designed or sold for incorporation into a dwelling. In determining whether a product is a consumer product, doubtful cases shall be resolved in favor of coverage. For a particular product received by a particular user, “normally used” refers to a typical or common use of that class of product, regardless of the status of the particular user or of the way in which the particular user actually uses, or expects or is expected to use, the product. A product is a consumer product regardless of whether the product has substantial commercial, industrial or non-consumer uses, unless such uses represent the only significant mode of use of the product.

“Installation Information” for a User Product means any methods, procedures, authorization keys, or other information required to install and execute modified versions of a covered work in that User Product from a modified version of its Corresponding Source. The information must suffice to ensure that the continued functioning of the modified object code is in no case prevented or interfered with solely because modification has been made.

If you convey an object code work under this section in, or with, or specifically for use in, a User Product, and the conveying occurs as part of a transaction in which the right of possession and use of the User Product is transferred to the recipient in perpetuity or for a fixed term (regardless of how the transaction is characterized), the Corresponding Source conveyed under this section must be accompanied by the Installation Information. But this requirement does not apply if neither you nor any third party retains the ability to install modified object code on the User Product (for example, the work has been installed in ROM).

The requirement to provide Installation Information does not include a requirement to continue to provide support service, warranty, or updates for a work that has been modified or installed by the recipient, or for the User Product in which it has been modified or installed. Access to a network may be denied when the modification itself materially and adversely affects the operation of the network or violates the rules and protocols for communication across the network.

Corresponding Source conveyed, and Installation Information provided, in accord with this section must be in a format that is publicly documented (and with an implementation available to the public in source code form), and must require no special password or key for unpacking, reading or copying.
7. Additional Terms.

“Additional permissions” are terms that supplement the terms of this License by making exceptions from one or more of its conditions. Additional permissions that are applicable to the entire Program shall be treated as though they were included in this License, to the extent that they are valid under applicable law. If additional permissions apply only to part of the Program, that part may be used separately under those permissions, but the entire Program remains governed by this License without regard to the additional permissions.

When you convey a copy of a covered work, you may at your option remove any additional permissions from that copy, or from any part of it. (Additional permissions may be written to require their own removal in certain cases when you modify the work.) You may place additional permissions on material, added by you to a covered work, for which you have or can give appropriate copyright permission.

Notwithstanding any other provision of this License, for material you add to a covered work, you may (if authorized by the copyright holders of that material) supplement the terms of this License with terms:

    a) Disclaiming warranty or limiting liability differently from the terms of sections 15 and 16 of this License; or
    b) Requiring preservation of specified reasonable legal notices or author attributions in that material or in the Appropriate Legal Notices displayed by works containing it; or
    c) Prohibiting misrepresentation of the origin of that material, or requiring that modified versions of such material be marked in reasonable ways as different from the original version; or
    d) Limiting the use for publicity purposes of names of licensors or authors of the material; or
    e) Declining to grant rights under trademark law for use of some trade names, trademarks, or service marks; or
    f) Requiring indemnification of licensors and authors of that material by anyone who conveys the material (or modified versions of it) with contractual assumptions of liability to the recipient, for any liability that these contractual assumptions directly impose on those licensors and authors.

All other non-permissive additional terms are considered “further restrictions” within the meaning of section 10. If the Program as you received it, or any part of it, contains a notice stating that it is governed by this License along with a term that is a further restriction, you may remove that term. If a license document contains a further restriction but permits relicensing or conveying under this License, you may add to a covered work material governed by the terms of that license document, provided that the further restriction does not survive such relicensing or conveying.

If you add terms to a covered work in accord with this section, you must place, in the relevant source files, a statement of the additional terms that apply to those files, or a notice indicating where to find the applicable terms.

Additional terms, permissive or non-permissive, may be stated in the form of a separately written license, or stated as exceptions; the above requirements apply either way.
8. Termination.

You may not propagate or modify a covered work except as expressly provided under this License. Any attempt otherwise to propagate or modify it is void, and will automatically terminate your rights under this License (including any patent licenses granted under the third paragraph of section 11).

However, if you cease all violation of this License, then your license from a particular copyright holder is reinstated (a) provisionally, unless and until the copyright holder explicitly and finally terminates your license, and (b) permanently, if the copyright holder fails to notify you of the violation by some reasonable means prior to 60 days after the cessation.

Moreover, your license from a particular copyright holder is reinstated permanently if the copyright holder notifies you of the violation by some reasonable means, this is the first time you have received notice of violation of this License (for any work) from that copyright holder, and you cure the violation prior to 30 days after your receipt of the notice.

Termination of your rights under this section does not terminate the licenses of parties who have received copies or rights from you under this License. If your rights have been terminated and not permanently reinstated, you do not qualify to receive new licenses for the same material under section 10.
9. Acceptance Not Required for Having Copies.

You are not required to accept this License in order to receive or run a copy of the Program. Ancillary propagation of a covered work occurring solely as a consequence of using peer-to-peer transmission to receive a copy likewise does not require acceptance. However, nothing other than this License grants you permission to propagate or modify any covered work. These actions infringe copyright if you do not accept this License. Therefore, by modifying or propagating a covered work, you indicate your acceptance of this License to do so.
10. Automatic Licensing of Downstream Recipients.

Each time you convey a covered work, the recipient automatically receives a license from the original licensors, to run, modify and propagate that work, subject to this License. You are not responsible for enforcing compliance by third parties with this License.

An “entity transaction” is a transaction transferring control of an organization, or substantially all assets of one, or subdividing an organization, or merging organizations. If propagation of a covered work results from an entity transaction, each party to that transaction who receives a copy of the work also receives whatever licenses to the work the party's predecessor in interest had or could give under the previous paragraph, plus a right to possession of the Corresponding Source of the work from the predecessor in interest, if the predecessor has it or can get it with reasonable efforts.

You may not impose any further restrictions on the exercise of the rights granted or affirmed under this License. For example, you may not impose a license fee, royalty, or other charge for exercise of rights granted under this License, and you may not initiate litigation (including a cross-claim or counterclaim in a lawsuit) alleging that any patent claim is infringed by making, using, selling, offering for sale, or importing the Program or any portion of it.
11. Patents.

A “contributor” is a copyright holder who authorizes use under this License of the Program or a work on which the Program is based. The work thus licensed is called the contributor's “contributor version”.

A contributor's “essential patent claims” are all patent claims owned or controlled by the contributor, whether already acquired or hereafter acquired, that would be infringed by some manner, permitted by this License, of making, using, or selling its contributor version, but do not include claims that would be infringed only as a consequence of further modification of the contributor version. For purposes of this definition, “control” includes the right to grant patent sublicenses in a manner consistent with the requirements of this License.

Each contributor grants you a non-exclusive, worldwide, royalty-free patent license under the contributor's essential patent claims, to make, use, sell, offer for sale, import and otherwise run, modify and propagate the contents of its contributor version.

In the following three paragraphs, a “patent license” is any express agreement or commitment, however denominated, not to enforce a patent (such as an express permission to practice a patent or covenant not to sue for patent infringement). To “grant” such a patent license to a party means to make such an agreement or commitment not to enforce a patent against the party.

If you convey a covered work, knowingly relying on a patent license, and the Corresponding Source of the work is not available for anyone to copy, free of charge and under the terms of this License, through a publicly available network server or other readily accessible means, then you must either (1) cause the Corresponding Source to be so available, or (2) arrange to deprive yourself of the benefit of the patent license for this particular work, or (3) arrange, in a manner consistent with the requirements of this License, to extend the patent license to downstream recipients. “Knowingly relying” means you have actual knowledge that, but for the patent license, your conveying the covered work in a country, or your recipient's use of the covered work in a country, would infringe one or more identifiable patents in that country that you have reason to believe are valid.

If, pursuant to or in connection with a single transaction or arrangement, you convey, or propagate by procuring conveyance of, a covered work, and grant a patent license to some of the parties receiving the covered work authorizing them to use, propagate, modify or convey a specific copy of the covered work, then the patent license you grant is automatically extended to all recipients of the covered work and works based on it.

A patent license is “discriminatory” if it does not include within the scope of its coverage, prohibits the exercise of, or is conditioned on the non-exercise of one or more of the rights that are specifically granted under this License. You may not convey a covered work if you are a party to an arrangement with a third party that is in the business of distributing software, under which you make payment to the third party based on the extent of your activity of conveying the work, and under which the third party grants, to any of the parties who would receive the covered work from you, a discriminatory patent license (a) in connection with copies of the covered work conveyed by you (or copies made from those copies), or (b) primarily for and in connection with specific products or compilations that contain the covered work, unless you entered into that arrangement, or that patent license was granted, prior to 28 March 2007.

Nothing in this License shall be construed as excluding or limiting any implied license or other defenses to infringement that may otherwise be available to you under applicable patent law.
12. No Surrender of Others' Freedom.

If conditions are imposed on you (whether by court order, agreement or otherwise) that contradict the conditions of this License, they do not excuse you from the conditions of this License. If you cannot convey a covered work so as to satisfy simultaneously your obligations under this License and any other pertinent obligations, then as a consequence you may not convey it at all. For example, if you agree to terms that obligate you to collect a royalty for further conveying from those to whom you convey the Program, the only way you could satisfy both those terms and this License would be to refrain entirely from conveying the Program.
13. Use with the GNU Affero General Public License.

Notwithstanding any other provision of this License, you have permission to link or combine any covered work with a work licensed under version 3 of the GNU Affero General Public License into a single combined work, and to convey the resulting work. The terms of this License will continue to apply to the part which is the covered work, but the special requirements of the GNU Affero General Public License, section 13, concerning interaction through a network will apply to the combination as such.
14. Revised Versions of this License.

The Free Software Foundation may publish revised and/or new versions of the GNU General Public License from time to time. Such new versions will be similar in spirit to the present version, but may differ in detail to address new problems or concerns.

Each version is given a distinguishing version number. If the Program specifies that a certain numbered version of the GNU General Public License “or any later version” applies to it, you have the option of following the terms and conditions either of that numbered version or of any later version published by the Free Software Foundation. If the Program does not specify a version number of the GNU General Public License, you may choose any version ever published by the Free Software Foundation.

If the Program specifies that a proxy can decide which future versions of the GNU General Public License can be used, that proxy's public statement of acceptance of a version permanently authorizes you to choose that version for the Program.

Later license versions may give you additional or different permissions. However, no additional obligations are imposed on any author or copyright holder as a result of your choosing to follow a later version.
15. Disclaimer of Warranty.

THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY APPLICABLE LAW. EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM “AS IS” WITHOUT WARRANTY OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE. THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM IS WITH YOU. SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF ALL NECESSARY SERVICING, REPAIR OR CORRECTION.
16. Limitation of Liability.

IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS), EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGES.
17. Interpretation of Sections 15 and 16.

If the disclaimer of warranty and limitation of liability provided above cannot be given local legal effect according to their terms, reviewing courts shall apply local law that most closely approximates an absolute waiver of all civil liability in connection with the Program, unless a warranty or assumption of liability accompanies a copy of the Program in return for a fee.

END OF TERMS AND CONDITIONS

================================================
FILE: LICENSE/LICENSE-AUTOENCODER
================================================
                                 Apache License
                           Version 2.0, January 2004
                        http://www.apache.org/licenses/

   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION

   1. Definitions.

      "License" shall mean the terms and conditions for use, reproduction,
      and distribution as defined by Sections 1 through 9 of this document.

      "Licensor" shall mean the copyright owner or entity authorized by
      the copyright owner that is granting the License.

      "Legal Entity" shall mean the union of the acting entity and all
      other entities that control, are controlled by, or are under common
      control with that entity. For the purposes of this definition,
      "control" means (i) the power, direct or indirect, to cause the
      direction or management of such entity, whether by contract or
      otherwise, or (ii) ownership of fifty percent (50%) or more of the
      outstanding shares, or (iii) beneficial ownership of such entity.

      "You" (or "Your") shall mean an individual or Legal Entity
      exercising permissions granted by this License.

      "Source" form shall mean the preferred form for making modifications,
      including but not limited to software source code, documentation
      source, and configuration files.

      "Object" form shall mean any form resulting from mechanical
      transformation or translation of a Source form, including but
      not limited to compiled object code, generated documentation,
      and conversions to other media types.

      "Work" shall mean the work of authorship, whether in Source or
      Object form, made available under the License, as indicated by a
      copyright notice that is included in or attached to the work
      (an example is provided in the Appendix below).

      "Derivative Works" shall mean any work, whether in Source or Object
      form, that is based on (or derived from) the Work and for which the
      editorial revisions, annotations, elaborations, or other modifications
      represent, as a whole, an original work of authorship. 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 and Derivative Works thereof.

      "Contribution" shall mean any work of authorship, including
      the original version of the Work and any modifications or additions
      to that Work or Derivative Works thereof, that is intentionally
      submitted to Licensor for inclusion in the Work by the copyright owner
      or by an individual or Legal Entity authorized to submit on behalf of
      the copyright owner. For the purposes of this definition, "submitted"
      means any form of electronic, verbal, or written communication sent
      to the Licensor or its representatives, including but not limited to
      communication on electronic mailing lists, source code control systems,
      and issue tracking systems that are managed by, or on behalf of, the
      Licensor for the purpose of discussing and improving the Work, but
      excluding communication that is conspicuously marked or otherwise
      designated in writing by the copyright owner as "Not a Contribution."

      "Contributor" shall mean Licensor and any individual or Legal Entity
      on behalf of whom a Contribution has been received by Licensor and
      subsequently incorporated within the Work.

   2. Grant of Copyright License. Subject to the terms and conditions of
      this License, each Contributor hereby grants to You a perpetual,
      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
      copyright license to reproduce, prepare Derivative Works of,
      publicly display, publicly perform, sublicense, and distribute the
      Work and such Derivative Works in Source or Object form.

   3. Grant of Patent License. Subject to the terms and conditions of
      this License, each Contributor hereby grants to You a perpetual,
      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
      (except as stated in this section) patent license to make, have made,
      use, offer to sell, sell, import, and otherwise transfer the Work,
      where such license applies only to those patent claims licensable
      by such Contributor that are necessarily infringed by their
      Contribution(s) alone or by combination of their Contribution(s)
      with the Work to which such Contribution(s) was submitted. If You
      institute patent litigation against any entity (including a
      cross-claim or counterclaim in a lawsuit) alleging that the Work
      or a Contribution incorporated within the Work constitutes direct
      or contributory patent infringement, then any patent licenses
      granted to You under this License for that Work shall terminate
      as of the date such litigation is filed.

   4. Redistribution. You may reproduce and distribute copies of the
      Work or Derivative Works thereof in any medium, with or without
      modifications, and in Source or Object form, provided that You
      meet the following conditions:

      (a) You must give any other recipients of the Work or
          Derivative Works a copy of this License; and

      (b) You must cause any modified files to carry prominent notices
          stating that You changed the files; and

      (c) You must retain, in the Source form of any Derivative Works
          that You distribute, all copyright, patent, trademark, and
          attribution notices from the Source form of the Work,
          excluding those notices that do not pertain to any part of
          the Derivative Works; and

      (d) If the Work includes a "NOTICE" text file as part of its
          distribution, then any Derivative Works that You distribute must
          include a readable copy of the attribution notices contained
          within such NOTICE file, excluding those notices that do not
          pertain to any part of the Derivative Works, in at least one
          of the following places: within a NOTICE text file distributed
          as part of the Derivative Works; within the Source form or
          documentation, if provided along with the Derivative Works; or,
          within a display generated by the Derivative Works, if and
          wherever such third-party notices normally appear. The contents
          of the NOTICE file are for informational purposes only and
          do not modify the License. You may add Your own attribution
          notices within Derivative Works that You distribute, alongside
          or as an addendum to the NOTICE text from the Work, provided
          that such additional attribution notices cannot be construed
          as modifying the License.

      You may add Your own copyright statement to Your modifications and
      may provide additional or different license terms and conditions
      for use, reproduction, or distribution of Your modifications, or
      for any such Derivative Works as a whole, provided Your use,
      reproduction, and distribution of the Work otherwise complies with
      the conditions stated in this License.

   5. Submission of Contributions. Unless You explicitly state otherwise,
      any Contribution intentionally submitted for inclusion in the Work
      by You to the Licensor shall be under the terms and conditions of
      this License, without any additional terms or conditions.
      Notwithstanding the above, nothing herein shall supersede or modify
      the terms of any separate license agreement you may have executed
      with Licensor regarding such Contributions.

   6. Trademarks. This License does not grant permission to use the trade
      names, trademarks, service marks, or product names of the Licensor,
      except as required for reasonable and customary use in describing the
      origin of the Work and reproducing the content of the NOTICE file.

   7. Disclaimer of Warranty. Unless required by applicable law or
      agreed to in writing, Licensor provides the Work (and each
      Contributor provides its Contributions) on an "AS IS" BASIS,
      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
      implied, including, without limitation, any warranties or conditions
      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
      PARTICULAR PURPOSE. You are solely responsible for determining the
      appropriateness of using or redistributing the Work and assume any
      risks associated with Your exercise of permissions under this License.

   8. Limitation of Liability. In no event and under no legal theory,
      whether in tort (including negligence), contract, or otherwise,
      unless required by applicable law (such as deliberate and grossly
      negligent acts) or agreed to in writing, shall any Contributor be
      liable to You for damages, including any direct, indirect, special,
      incidental, or consequential damages of any character arising as a
      result of this License or out of the use or inability to use the
      Work (including but not limited to damages for loss of goodwill,
      work stoppage, computer failure or malfunction, or any and all
      other commercial damages or losses), even if such Contributor
      has been advised of the possibility of such damages.

   9. Accepting Warranty or Additional Liability. While redistributing
      the Work or Derivative Works thereof, You may choose to offer,
      and charge a fee for, acceptance of support, warranty, indemnity,
      or other liability obligations and/or rights consistent with this
      License. However, in accepting such obligations, You may act only
      on Your own behalf and on Your sole responsibility, not on behalf
      of any other Contributor, and only if You agree to indemnify,
      defend, and hold each Contributor harmless for any liability
      incurred by, or claims asserted against, such Contributor by reason
      of your accepting any such warranty or additional liability.

   END OF TERMS AND CONDITIONS

   APPENDIX: How to apply the Apache License to your work.

      To apply the Apache License to your work, attach the following
      boilerplate notice, with the fields enclosed by brackets "[]"
      replaced with your own identifying information. (Don't include
      the brackets!)  The text should be enclosed in the appropriate
      comment syntax for the file format. We also recommend that a
      file or class name and description of purpose be included on the
      same "printed page" as the copyright notice for easier
      identification within third-party archives.

   Copyright [yyyy] [name of copyright owner]

   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.


================================================
FILE: LICENSE/LICENSE-CONTRASTIVE-LEARNER
================================================
MIT License

Copyright (c) 2020 Phil Wang

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.


================================================
FILE: LICENSE/LICENSE-FID
================================================
                                 Apache License
                           Version 2.0, January 2004
                        http://www.apache.org/licenses/

   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION

   1. Definitions.

      "License" shall mean the terms and conditions for use, reproduction,
      and distribution as defined by Sections 1 through 9 of this document.

      "Licensor" shall mean the copyright owner or entity authorized by
      the copyright owner that is granting the License.

      "Legal Entity" shall mean the union of the acting entity and all
      other entities that control, are controlled by, or are under common
      control with that entity. For the purposes of this definition,
      "control" means (i) the power, direct or indirect, to cause the
      direction or management of such entity, whether by contract or
      otherwise, or (ii) ownership of fifty percent (50%) or more of the
      outstanding shares, or (iii) beneficial ownership of such entity.

      "You" (or "Your") shall mean an individual or Legal Entity
      exercising permissions granted by this License.

      "Source" form shall mean the preferred form for making modifications,
      including but not limited to software source code, documentation
      source, and configuration files.

      "Object" form shall mean any form resulting from mechanical
      transformation or translation of a Source form, including but
      not limited to compiled object code, generated documentation,
      and conversions to other media types.

      "Work" shall mean the work of authorship, whether in Source or
      Object form, made available under the License, as indicated by a
      copyright notice that is included in or attached to the work
      (an example is provided in the Appendix below).

      "Derivative Works" shall mean any work, whether in Source or Object
      form, that is based on (or derived from) the Work and for which the
      editorial revisions, annotations, elaborations, or other modifications
      represent, as a whole, an original work of authorship. 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 and Derivative Works thereof.

      "Contribution" shall mean any work of authorship, including
      the original version of the Work and any modifications or additions
      to that Work or Derivative Works thereof, that is intentionally
      submitted to Licensor for inclusion in the Work by the copyright owner
      or by an individual or Legal Entity authorized to submit on behalf of
      the copyright owner. For the purposes of this definition, "submitted"
      means any form of electronic, verbal, or written communication sent
      to the Licensor or its representatives, including but not limited to
      communication on electronic mailing lists, source code control systems,
      and issue tracking systems that are managed by, or on behalf of, the
      Licensor for the purpose of discussing and improving the Work, but
      excluding communication that is conspicuously marked or otherwise
      designated in writing by the copyright owner as "Not a Contribution."

      "Contributor" shall mean Licensor and any individual or Legal Entity
      on behalf of whom a Contribution has been received by Licensor and
      subsequently incorporated within the Work.

   2. Grant of Copyright License. Subject to the terms and conditions of
      this License, each Contributor hereby grants to You a perpetual,
      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
      copyright license to reproduce, prepare Derivative Works of,
      publicly display, publicly perform, sublicense, and distribute the
      Work and such Derivative Works in Source or Object form.

   3. Grant of Patent License. Subject to the terms and conditions of
      this License, each Contributor hereby grants to You a perpetual,
      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
      (except as stated in this section) patent license to make, have made,
      use, offer to sell, sell, import, and otherwise transfer the Work,
      where such license applies only to those patent claims licensable
      by such Contributor that are necessarily infringed by their
      Contribution(s) alone or by combination of their Contribution(s)
      with the Work to which such Contribution(s) was submitted. If You
      institute patent litigation against any entity (including a
      cross-claim or counterclaim in a lawsuit) alleging that the Work
      or a Contribution incorporated within the Work constitutes direct
      or contributory patent infringement, then any patent licenses
      granted to You under this License for that Work shall terminate
      as of the date such litigation is filed.

   4. Redistribution. You may reproduce and distribute copies of the
      Work or Derivative Works thereof in any medium, with or without
      modifications, and in Source or Object form, provided that You
      meet the following conditions:

      (a) You must give any other recipients of the Work or
          Derivative Works a copy of this License; and

      (b) You must cause any modified files to carry prominent notices
          stating that You changed the files; and

      (c) You must retain, in the Source form of any Derivative Works
          that You distribute, all copyright, patent, trademark, and
          attribution notices from the Source form of the Work,
          excluding those notices that do not pertain to any part of
          the Derivative Works; and

      (d) If the Work includes a "NOTICE" text file as part of its
          distribution, then any Derivative Works that You distribute must
          include a readable copy of the attribution notices contained
          within such NOTICE file, excluding those notices that do not
          pertain to any part of the Derivative Works, in at least one
          of the following places: within a NOTICE text file distributed
          as part of the Derivative Works; within the Source form or
          documentation, if provided along with the Derivative Works; or,
          within a display generated by the Derivative Works, if and
          wherever such third-party notices normally appear. The contents
          of the NOTICE file are for informational purposes only and
          do not modify the License. You may add Your own attribution
          notices within Derivative Works that You distribute, alongside
          or as an addendum to the NOTICE text from the Work, provided
          that such additional attribution notices cannot be construed
          as modifying the License.

      You may add Your own copyright statement to Your modifications and
      may provide additional or different license terms and conditions
      for use, reproduction, or distribution of Your modifications, or
      for any such Derivative Works as a whole, provided Your use,
      reproduction, and distribution of the Work otherwise complies with
      the conditions stated in this License.

   5. Submission of Contributions. Unless You explicitly state otherwise,
      any Contribution intentionally submitted for inclusion in the Work
      by You to the Licensor shall be under the terms and conditions of
      this License, without any additional terms or conditions.
      Notwithstanding the above, nothing herein shall supersede or modify
      the terms of any separate license agreement you may have executed
      with Licensor regarding such Contributions.

   6. Trademarks. This License does not grant permission to use the trade
      names, trademarks, service marks, or product names of the Licensor,
      except as required for reasonable and customary use in describing the
      origin of the Work and reproducing the content of the NOTICE file.

   7. Disclaimer of Warranty. Unless required by applicable law or
      agreed to in writing, Licensor provides the Work (and each
      Contributor provides its Contributions) on an "AS IS" BASIS,
      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
      implied, including, without limitation, any warranties or conditions
      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
      PARTICULAR PURPOSE. You are solely responsible for determining the
      appropriateness of using or redistributing the Work and assume any
      risks associated with Your exercise of permissions under this License.

   8. Limitation of Liability. In no event and under no legal theory,
      whether in tort (including negligence), contract, or otherwise,
      unless required by applicable law (such as deliberate and grossly
      negligent acts) or agreed to in writing, shall any Contributor be
      liable to You for damages, including any direct, indirect, special,
      incidental, or consequential damages of any character arising as a
      result of this License or out of the use or inability to use the
      Work (including but not limited to damages for loss of goodwill,
      work stoppage, computer failure or malfunction, or any and all
      other commercial damages or losses), even if such Contributor
      has been advised of the possibility of such damages.

   9. Accepting Warranty or Additional Liability. While redistributing
      the Work or Derivative Works thereof, You may choose to offer,
      and charge a fee for, acceptance of support, warranty, indemnity,
      or other liability obligations and/or rights consistent with this
      License. However, in accepting such obligations, You may act only
      on Your own behalf and on Your sole responsibility, not on behalf
      of any other Contributor, and only if You agree to indemnify,
      defend, and hold each Contributor harmless for any liability
      incurred by, or claims asserted against, such Contributor by reason
      of your accepting any such warranty or additional liability.

   END OF TERMS AND CONDITIONS

   APPENDIX: How to apply the Apache License to your work.

      To apply the Apache License to your work, attach the following
      boilerplate notice, with the fields enclosed by brackets "[]"
      replaced with your own identifying information. (Don't include
      the brackets!)  The text should be enclosed in the appropriate
      comment syntax for the file format. We also recommend that a
      file or class name and description of purpose be included on the
      same "printed page" as the copyright notice for easier
      identification within third-party archives.

   Copyright [yyyy] [name of copyright owner]

   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.


================================================
FILE: LICENSE/LICENSE-LPIPS
================================================
Copyright (c) 2018, Richard Zhang, Phillip Isola, Alexei A. Efros, Eli Shechtman, Oliver Wang
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

* Redistributions of source code must retain the above copyright notice, this
  list of conditions and the following disclaimer.

* Redistributions in binary form must reproduce the above copyright notice,
  this list of conditions and the following disclaimer in the documentation
  and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.



================================================
FILE: LICENSE/LICENSE-LUCIDRAINS
================================================
                    GNU GENERAL PUBLIC LICENSE
                       Version 3, 29 June 2007

 Copyright (C) 2007 Free Software Foundation, Inc. <https://fsf.org/>
 Everyone is permitted to copy and distribute verbatim copies
 of this license document, but changing it is not allowed.

                            Preamble

  The GNU General Public License is a free, copyleft license for
software and other kinds of works.

  The licenses for most software and other practical works are designed
to take away your freedom to share and change the works.  By contrast,
the GNU General Public License is intended to guarantee your freedom to
share and change all versions of a program--to make sure it remains free
software for all its users.  We, the Free Software Foundation, use the
GNU General Public License for most of our software; it applies also to
any other work released this way by its authors.  You can apply it to
your programs, too.

  When we speak of free software, we are referring to freedom, not
price.  Our General Public Licenses are designed to make sure that you
have the freedom to distribute copies of free software (and charge for
them if you wish), that you receive source code or can get it if you
want it, that you can change the software or use pieces of it in new
free programs, and that you know you can do these things.

  To protect your rights, we need to prevent others from denying you
these rights or asking you to surrender the rights.  Therefore, you have
certain responsibilities if you distribute copies of the software, or if
you modify it: responsibilities to respect the freedom of others.

  For example, if you distribute copies of such a program, whether
gratis or for a fee, you must pass on to the recipients the same
freedoms that you received.  You must make sure that they, too, receive
or can get the source code.  And you must show them these terms so they
know their rights.

  Developers that use the GNU GPL protect your rights with two steps:
(1) assert copyright on the software, and (2) offer you this License
giving you legal permission to copy, distribute and/or modify it.

  For the developers' and authors' protection, the GPL clearly explains
that there is no warranty for this free software.  For both users' and
authors' sake, the GPL requires that modified versions be marked as
changed, so that their problems will not be attributed erroneously to
authors of previous versions.

  Some devices are designed to deny users access to install or run
modified versions of the software inside them, although the manufacturer
can do so.  This is fundamentally incompatible with the aim of
protecting users' freedom to change the software.  The systematic
pattern of such abuse occurs in the area of products for individuals to
use, which is precisely where it is most unacceptable.  Therefore, we
have designed this version of the GPL to prohibit the practice for those
products.  If such problems arise substantially in other domains, we
stand ready to extend this provision to those domains in future versions
of the GPL, as needed to protect the freedom of users.

  Finally, every program is threatened constantly by software patents.
States should not allow patents to restrict development and use of
software on general-purpose computers, but in those that do, we wish to
avoid the special danger that patents applied to a free program could
make it effectively proprietary.  To prevent this, the GPL assures that
patents cannot be used to render the program non-free.

  The precise terms and conditions for copying, distribution and
modification follow.

                       TERMS AND CONDITIONS

  0. Definitions.

  "This License" refers to version 3 of the GNU General Public License.

  "Copyright" also means copyright-like laws that apply to other kinds of
works, such as semiconductor masks.

  "The Program" refers to any copyrightable work licensed under this
License.  Each licensee is addressed as "you".  "Licensees" and
"recipients" may be individuals or organizations.

  To "modify" a work means to copy from or adapt all or part of the work
in a fashion requiring copyright permission, other than the making of an
exact copy.  The resulting work is called a "modified version" of the
earlier work or a work "based on" the earlier work.

  A "covered work" means either the unmodified Program or a work based
on the Program.

  To "propagate" a work means to do anything with it that, without
permission, would make you directly or secondarily liable for
infringement under applicable copyright law, except executing it on a
computer or modifying a private copy.  Propagation includes copying,
distribution (with or without modification), making available to the
public, and in some countries other activities as well.

  To "convey" a work means any kind of propagation that enables other
parties to make or receive copies.  Mere interaction with a user through
a computer network, with no transfer of a copy, is not conveying.

  An interactive user interface displays "Appropriate Legal Notices"
to the extent that it includes a convenient and prominently visible
feature that (1) displays an appropriate copyright notice, and (2)
tells the user that there is no warranty for the work (except to the
extent that warranties are provided), that licensees may convey the
work under this License, and how to view a copy of this License.  If
the interface presents a list of user commands or options, such as a
menu, a prominent item in the list meets this criterion.

  1. Source Code.

  The "source code" for a work means the preferred form of the work
for making modifications to it.  "Object code" means any non-source
form of a work.

  A "Standard Interface" means an interface that either is an official
standard defined by a recognized standards body, or, in the case of
interfaces specified for a particular programming language, one that
is widely used among developers working in that language.

  The "System Libraries" of an executable work include anything, other
than the work as a whole, that (a) is included in the normal form of
packaging a Major Component, but which is not part of that Major
Component, and (b) serves only to enable use of the work with that
Major Component, or to implement a Standard Interface for which an
implementation is available to the public in source code form.  A
"Major Component", in this context, means a major essential component
(kernel, window system, and so on) of the specific operating system
(if any) on which the executable work runs, or a compiler used to
produce the work, or an object code interpreter used to run it.

  The "Corresponding Source" for a work in object code form means all
the source code needed to generate, install, and (for an executable
work) run the object code and to modify the work, including scripts to
control those activities.  However, it does not include the work's
System Libraries, or general-purpose tools or generally available free
programs which are used unmodified in performing those activities but
which are not part of the work.  For example, Corresponding Source
includes interface definition files associated with source files for
the work, and the source code for shared libraries and dynamically
linked subprograms that the work is specifically designed to require,
such as by intimate data communication or control flow between those
subprograms and other parts of the work.

  The Corresponding Source need not include anything that users
can regenerate automatically from other parts of the Corresponding
Source.

  The Corresponding Source for a work in source code form is that
same work.

  2. Basic Permissions.

  All rights granted under this License are granted for the term of
copyright on the Program, and are irrevocable provided the stated
conditions are met.  This License explicitly affirms your unlimited
permission to run the unmodified Program.  The output from running a
covered work is covered by this License only if the output, given its
content, constitutes a covered work.  This License acknowledges your
rights of fair use or other equivalent, as provided by copyright law.

  You may make, run and propagate covered works that you do not
convey, without conditions so long as your license otherwise remains
in force.  You may convey covered works to others for the sole purpose
of having them make modifications exclusively for you, or provide you
with facilities for running those works, provided that you comply with
the terms of this License in conveying all material for which you do
not control copyright.  Those thus making or running the covered works
for you must do so exclusively on your behalf, under your direction
and control, on terms that prohibit them from making any copies of
your copyrighted material outside their relationship with you.

  Conveying under any other circumstances is permitted solely under
the conditions stated below.  Sublicensing is not allowed; section 10
makes it unnecessary.

  3. Protecting Users' Legal Rights From Anti-Circumvention Law.

  No covered work shall be deemed part of an effective technological
measure under any applicable law fulfilling obligations under article
11 of the WIPO copyright treaty adopted on 20 December 1996, or
similar laws prohibiting or restricting circumvention of such
measures.

  When you convey a covered work, you waive any legal power to forbid
circumvention of technological measures to the extent such circumvention
is effected by exercising rights under this License with respect to
the covered work, and you disclaim any intention to limit operation or
modification of the work as a means of enforcing, against the work's
users, your or third parties' legal rights to forbid circumvention of
technological measures.

  4. Conveying Verbatim Copies.

  You may convey verbatim copies of the Program's source code as you
receive it, in any medium, provided that you conspicuously and
appropriately publish on each copy an appropriate copyright notice;
keep intact all notices stating that this License and any
non-permissive terms added in accord with section 7 apply to the code;
keep intact all notices of the absence of any warranty; and give all
recipients a copy of this License along with the Program.

  You may charge any price or no price for each copy that you convey,
and you may offer support or warranty protection for a fee.

  5. Conveying Modified Source Versions.

  You may convey a work based on the Program, or the modifications to
produce it from the Program, in the form of source code under the
terms of section 4, provided that you also meet all of these conditions:

    a) The work must carry prominent notices stating that you modified
    it, and giving a relevant date.

    b) The work must carry prominent notices stating that it is
    released under this License and any conditions added under section
    7.  This requirement modifies the requirement in section 4 to
    "keep intact all notices".

    c) You must license the entire work, as a whole, under this
    License to anyone who comes into possession of a copy.  This
    License will therefore apply, along with any applicable section 7
    additional terms, to the whole of the work, and all its parts,
    regardless of how they are packaged.  This License gives no
    permission to license the work in any other way, but it does not
    invalidate such permission if you have separately received it.

    d) If the work has interactive user interfaces, each must display
    Appropriate Legal Notices; however, if the Program has interactive
    interfaces that do not display Appropriate Legal Notices, your
    work need not make them do so.

  A compilation of a covered work with other separate and independent
works, which are not by their nature extensions of the covered work,
and which are not combined with it such as to form a larger program,
in or on a volume of a storage or distribution medium, is called an
"aggregate" if the compilation and its resulting copyright are not
used to limit the access or legal rights of the compilation's users
beyond what the individual works permit.  Inclusion of a covered work
in an aggregate does not cause this License to apply to the other
parts of the aggregate.

  6. Conveying Non-Source Forms.

  You may convey a covered work in object code form under the terms
of sections 4 and 5, provided that you also convey the
machine-readable Corresponding Source under the terms of this License,
in one of these ways:

    a) Convey the object code in, or embodied in, a physical product
    (including a physical distribution medium), accompanied by the
    Corresponding Source fixed on a durable physical medium
    customarily used for software interchange.

    b) Convey the object code in, or embodied in, a physical product
    (including a physical distribution medium), accompanied by a
    written offer, valid for at least three years and valid for as
    long as you offer spare parts or customer support for that product
    model, to give anyone who possesses the object code either (1) a
    copy of the Corresponding Source for all the software in the
    product that is covered by this License, on a durable physical
    medium customarily used for software interchange, for a price no
    more than your reasonable cost of physically performing this
    conveying of source, or (2) access to copy the
    Corresponding Source from a network server at no charge.

    c) Convey individual copies of the object code with a copy of the
    written offer to provide the Corresponding Source.  This
    alternative is allowed only occasionally and noncommercially, and
    only if you received the object code with such an offer, in accord
    with subsection 6b.

    d) Convey the object code by offering access from a designated
    place (gratis or for a charge), and offer equivalent access to the
    Corresponding Source in the same way through the same place at no
    further charge.  You need not require recipients to copy the
    Corresponding Source along with the object code.  If the place to
    copy the object code is a network server, the Corresponding Source
    may be on a different server (operated by you or a third party)
    that supports equivalent copying facilities, provided you maintain
    clear directions next to the object code saying where to find the
    Corresponding Source.  Regardless of what server hosts the
    Corresponding Source, you remain obligated to ensure that it is
    available for as long as needed to satisfy these requirements.

    e) Convey the object code using peer-to-peer transmission, provided
    you inform other peers where the object code and Corresponding
    Source of the work are being offered to the general public at no
    charge under subsection 6d.

  A separable portion of the object code, whose source code is excluded
from the Corresponding Source as a System Library, need not be
included in conveying the object code work.

  A "User Product" is either (1) a "consumer product", which means any
tangible personal property which is normally used for personal, family,
or household purposes, or (2) anything designed or sold for incorporation
into a dwelling.  In determining whether a product is a consumer product,
doubtful cases shall be resolved in favor of coverage.  For a particular
product received by a particular user, "normally used" refers to a
typical or common use of that class of product, regardless of the status
of the particular user or of the way in which the particular user
actually uses, or expects or is expected to use, the product.  A product
is a consumer product regardless of whether the product has substantial
commercial, industrial or non-consumer uses, unless such uses represent
the only significant mode of use of the product.

  "Installation Information" for a User Product means any methods,
procedures, authorization keys, or other information required to install
and execute modified versions of a covered work in that User Product from
a modified version of its Corresponding Source.  The information must
suffice to ensure that the continued functioning of the modified object
code is in no case prevented or interfered with solely because
modification has been made.

  If you convey an object code work under this section in, or with, or
specifically for use in, a User Product, and the conveying occurs as
part of a transaction in which the right of possession and use of the
User Product is transferred to the recipient in perpetuity or for a
fixed term (regardless of how the transaction is characterized), the
Corresponding Source conveyed under this section must be accompanied
by the Installation Information.  But this requirement does not apply
if neither you nor any third party retains the ability to install
modified object code on the User Product (for example, the work has
been installed in ROM).

  The requirement to provide Installation Information does not include a
requirement to continue to provide support service, warranty, or updates
for a work that has been modified or installed by the recipient, or for
the User Product in which it has been modified or installed.  Access to a
network may be denied when the modification itself materially and
adversely affects the operation of the network or violates the rules and
protocols for communication across the network.

  Corresponding Source conveyed, and Installation Information provided,
in accord with this section must be in a format that is publicly
documented (and with an implementation available to the public in
source code form), and must require no special password or key for
unpacking, reading or copying.

  7. Additional Terms.

  "Additional permissions" are terms that supplement the terms of this
License by making exceptions from one or more of its conditions.
Additional permissions that are applicable to the entire Program shall
be treated as though they were included in this License, to the extent
that they are valid under applicable law.  If additional permissions
apply only to part of the Program, that part may be used separately
under those permissions, but the entire Program remains governed by
this License without regard to the additional permissions.

  When you convey a copy of a covered work, you may at your option
remove any additional permissions from that copy, or from any part of
it.  (Additional permissions may be written to require their own
removal in certain cases when you modify the work.)  You may place
additional permissions on material, added by you to a covered work,
for which you have or can give appropriate copyright permission.

  Notwithstanding any other provision of this License, for material you
add to a covered work, you may (if authorized by the copyright holders of
that material) supplement the terms of this License with terms:

    a) Disclaiming warranty or limiting liability differently from the
    terms of sections 15 and 16 of this License; or

    b) Requiring preservation of specified reasonable legal notices or
    author attributions in that material or in the Appropriate Legal
    Notices displayed by works containing it; or

    c) Prohibiting misrepresentation of the origin of that material, or
    requiring that modified versions of such material be marked in
    reasonable ways as different from the original version; or

    d) Limiting the use for publicity purposes of names of licensors or
    authors of the material; or

    e) Declining to grant rights under trademark law for use of some
    trade names, trademarks, or service marks; or

    f) Requiring indemnification of licensors and authors of that
    material by anyone who conveys the material (or modified versions of
    it) with contractual assumptions of liability to the recipient, for
    any liability that these contractual assumptions directly impose on
    those licensors and authors.

  All other non-permissive additional terms are considered "further
restrictions" within the meaning of section 10.  If the Program as you
received it, or any part of it, contains a notice stating that it is
governed by this License along with a term that is a further
restriction, you may remove that term.  If a license document contains
a further restriction but permits relicensing or conveying under this
License, you may add to a covered work material governed by the terms
of that license document, provided that the further restriction does
not survive such relicensing or conveying.

  If you add terms to a covered work in accord with this section, you
must place, in the relevant source files, a statement of the
additional terms that apply to those files, or a notice indicating
where to find the applicable terms.

  Additional terms, permissive or non-permissive, may be stated in the
form of a separately written license, or stated as exceptions;
the above requirements apply either way.

  8. Termination.

  You may not propagate or modify a covered work except as expressly
provided under this License.  Any attempt otherwise to propagate or
modify it is void, and will automatically terminate your rights under
this License (including any patent licenses granted under the third
paragraph of section 11).

  However, if you cease all violation of this License, then your
license from a particular copyright holder is reinstated (a)
provisionally, unless and until the copyright holder explicitly and
finally terminates your license, and (b) permanently, if the copyright
holder fails to notify you of the violation by some reasonable means
prior to 60 days after the cessation.

  Moreover, your license from a particular copyright holder is
reinstated permanently if the copyright holder notifies you of the
violation by some reasonable means, this is the first time you have
received notice of violation of this License (for any work) from that
copyright holder, and you cure the violation prior to 30 days after
your receipt of the notice.

  Termination of your rights under this section does not terminate the
licenses of parties who have received copies or rights from you under
this License.  If your rights have been terminated and not permanently
reinstated, you do not qualify to receive new licenses for the same
material under section 10.

  9. Acceptance Not Required for Having Copies.

  You are not required to accept this License in order to receive or
run a copy of the Program.  Ancillary propagation of a covered work
occurring solely as a consequence of using peer-to-peer transmission
to receive a copy likewise does not require acceptance.  However,
nothing other than this License grants you permission to propagate or
modify any covered work.  These actions infringe copyright if you do
not accept this License.  Therefore, by modifying or propagating a
covered work, you indicate your acceptance of this License to do so.

  10. Automatic Licensing of Downstream Recipients.

  Each time you convey a covered work, the recipient automatically
receives a license from the original licensors, to run, modify and
propagate that work, subject to this License.  You are not responsible
for enforcing compliance by third parties with this License.

  An "entity transaction" is a transaction transferring control of an
organization, or substantially all assets of one, or subdividing an
organization, or merging organizations.  If propagation of a covered
work results from an entity transaction, each party to that
transaction who receives a copy of the work also receives whatever
licenses to the work the party's predecessor in interest had or could
give under the previous paragraph, plus a right to possession of the
Corresponding Source of the work from the predecessor in interest, if
the predecessor has it or can get it with reasonable efforts.

  You may not impose any further restrictions on the exercise of the
rights granted or affirmed under this License.  For example, you may
not impose a license fee, royalty, or other charge for exercise of
rights granted under this License, and you may not initiate litigation
(including a cross-claim or counterclaim in a lawsuit) alleging that
any patent claim is infringed by making, using, selling, offering for
sale, or importing the Program or any portion of it.

  11. Patents.

  A "contributor" is a copyright holder who authorizes use under this
License of the Program or a work on which the Program is based.  The
work thus licensed is called the contributor's "contributor version".

  A contributor's "essential patent claims" are all patent claims
owned or controlled by the contributor, whether already acquired or
hereafter acquired, that would be infringed by some manner, permitted
by this License, of making, using, or selling its contributor version,
but do not include claims that would be infringed only as a
consequence of further modification of the contributor version.  For
purposes of this definition, "control" includes the right to grant
patent sublicenses in a manner consistent with the requirements of
this License.

  Each contributor grants you a non-exclusive, worldwide, royalty-free
patent license under the contributor's essential patent claims, to
make, use, sell, offer for sale, import and otherwise run, modify and
propagate the contents of its contributor version.

  In the following three paragraphs, a "patent license" is any express
agreement or commitment, however denominated, not to enforce a patent
(such as an express permission to practice a patent or covenant not to
sue for patent infringement).  To "grant" such a patent license to a
party means to make such an agreement or commitment not to enforce a
patent against the party.

  If you convey a covered work, knowingly relying on a patent license,
and the Corresponding Source of the work is not available for anyone
to copy, free of charge and under the terms of this License, through a
publicly available network server or other readily accessible means,
then you must either (1) cause the Corresponding Source to be so
available, or (2) arrange to deprive yourself of the benefit of the
patent license for this particular work, or (3) arrange, in a manner
consistent with the requirements of this License, to extend the patent
license to downstream recipients.  "Knowingly relying" means you have
actual knowledge that, but for the patent license, your conveying the
covered work in a country, or your recipient's use of the covered work
in a country, would infringe one or more identifiable patents in that
country that you have reason to believe are valid.

  If, pursuant to or in connection with a single transaction or
arrangement, you convey, or propagate by procuring conveyance of, a
covered work, and grant a patent license to some of the parties
receiving the covered work authorizing them to use, propagate, modify
or convey a specific copy of the covered work, then the patent license
you grant is automatically extended to all recipients of the covered
work and works based on it.

  A patent license is "discriminatory" if it does not include within
the scope of its coverage, prohibits the exercise of, or is
conditioned on the non-exercise of one or more of the rights that are
specifically granted under this License.  You may not convey a covered
work if you are a party to an arrangement with a third party that is
in the business of distributing software, under which you make payment
to the third party based on the extent of your activity of conveying
the work, and under which the third party grants, to any of the
parties who would receive the covered work from you, a discriminatory
patent license (a) in connection with copies of the covered work
conveyed by you (or copies made from those copies), or (b) primarily
for and in connection with specific products or compilations that
contain the covered work, unless you entered into that arrangement,
or that patent license was granted, prior to 28 March 2007.

  Nothing in this License shall be construed as excluding or limiting
any implied license or other defenses to infringement that may
otherwise be available to you under applicable patent law.

  12. No Surrender of Others' Freedom.

  If conditions are imposed on you (whether by court order, agreement or
otherwise) that contradict the conditions of this License, they do not
excuse you from the conditions of this License.  If you cannot convey a
covered work so as to satisfy simultaneously your obligations under this
License and any other pertinent obligations, then as a consequence you may
not convey it at all.  For example, if you agree to terms that obligate you
to collect a royalty for further conveying from those to whom you convey
the Program, the only way you could satisfy both those terms and this
License would be to refrain entirely from conveying the Program.

  13. Use with the GNU Affero General Public License.

  Notwithstanding any other provision of this License, you have
permission to link or combine any covered work with a work licensed
under version 3 of the GNU Affero General Public License into a single
combined work, and to convey the resulting work.  The terms of this
License will continue to apply to the part which is the covered work,
but the special requirements of the GNU Affero General Public License,
section 13, concerning interaction through a network will apply to the
combination as such.

  14. Revised Versions of this License.

  The Free Software Foundation may publish revised and/or new versions of
the GNU General Public License from time to time.  Such new versions will
be similar in spirit to the present version, but may differ in detail to
address new problems or concerns.

  Each version is given a distinguishing version number.  If the
Program specifies that a certain numbered version of the GNU General
Public License "or any later version" applies to it, you have the
option of following the terms and conditions either of that numbered
version or of any later version published by the Free Software
Foundation.  If the Program does not specify a version number of the
GNU General Public License, you may choose any version ever published
by the Free Software Foundation.

  If the Program specifies that a proxy can decide which future
versions of the GNU General Public License can be used, that proxy's
public statement of acceptance of a version permanently authorizes you
to choose that version for the Program.

  Later license versions may give you additional or different
permissions.  However, no additional obligations are imposed on any
author or copyright holder as a result of your choosing to follow a
later version.

  15. Disclaimer of Warranty.

  THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY
APPLICABLE LAW.  EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT
HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY
OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO,
THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
PURPOSE.  THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM
IS WITH YOU.  SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF
ALL NECESSARY SERVICING, REPAIR OR CORRECTION.

  16. Limitation of Liability.

  IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING
WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS
THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY
GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE
USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF
DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD
PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS),
EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF
SUCH DAMAGES.

  17. Interpretation of Sections 15 and 16.

  If the disclaimer of warranty and limitation of liability provided
above cannot be given local legal effect according to their terms,
reviewing courts shall apply local law that most closely approximates
an absolute waiver of all civil liability in connection with the
Program, unless a warranty or assumption of liability accompanies a
copy of the Program in return for a fee.

                     END OF TERMS AND CONDITIONS

            How to Apply These Terms to Your New Programs

  If you develop a new program, and you want it to be of the greatest
possible use to the public, the best way to achieve this is to make it
free software which everyone can redistribute and change under these terms.

  To do so, attach the following notices to the program.  It is safest
to attach them to the start of each source file to most effectively
state the exclusion of warranty; and each file should have at least
the "copyright" line and a pointer to where the full notice is found.

    <one line to give the program's name and a brief idea of what it does.>
    Copyright (C) <year>  <name of author>

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <https://www.gnu.org/licenses/>.

Also add information on how to contact you by electronic and paper mail.

  If the program does terminal interaction, make it output a short
notice like this when it starts in an interactive mode:

    <program>  Copyright (C) <year>  <name of author>
    This program comes with ABSOLUTELY NO WARRANTY; for details type `show w'.
    This is free software, and you are welcome to redistribute it
    under certain conditions; type `show c' for details.

The hypothetical commands `show w' and `show c' should show the appropriate
parts of the General Public License.  Of course, your program's commands
might be different; for a GUI interface, you would use an "about box".

  You should also get your employer (if you work as a programmer) or school,
if any, to sign a "copyright disclaimer" for the program, if necessary.
For more information on this, and how to apply and follow the GNU GPL, see
<https://www.gnu.org/licenses/>.

  The GNU General Public License does not permit incorporating your program
into proprietary programs.  If your program is a subroutine library, you
may consider it more useful to permit linking proprietary applications with
the library.  If this is what you want to do, use the GNU Lesser General
Public License instead of this License.  But first, please read
<https://www.gnu.org/licenses/why-not-lgpl.html>.


================================================
FILE: LICENSE/LICENSE-NVIDIA
================================================
Copyright (c) 2019, NVIDIA Corporation. All rights reserved.


Nvidia Source Code License-NC

=======================================================================

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.

"Nvidia Processors" means any central processing unit (CPU), graphics
processing unit (GPU), field-programmable gate array (FPGA),
application-specific integrated circuit (ASIC) or any combination
thereof designed, made, sold, or provided by Nvidia or its affiliates.

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. The Work or
    derivative works thereof may be used or intended for use by Nvidia
    or its affiliates commercially or non-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 grants in Sections 2.1 and 2.2) 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 grants in Sections 2.1 and
    2.2) 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: LICENSE/LICENSE-ROSINALITY
================================================
MIT License

Copyright (c) 2019 Kim Seonghyeon

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.


================================================
FILE: LICENSE/LICENSE-VGG
================================================
Copyright (C) 2019 NVIDIA Corporation. Ting-Chun Wang, Ming-Yu Liu, Jun-Yan Zhu.
BSD License. All rights reserved. 

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

* Redistributions of source code must retain the above copyright notice, this
  list of conditions and the following disclaimer.

* Redistributions in binary form must reproduce the above copyright notice,
  this list of conditions and the following disclaimer in the documentation
  and/or other materials provided with the distribution.

THE AUTHOR DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS SOFTWARE, INCLUDING ALL 
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR ANY PARTICULAR PURPOSE. 
IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, INDIRECT OR CONSEQUENTIAL 
DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, 
WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING 
OUT OF OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.


--------------------------- LICENSE FOR pytorch-CycleGAN-and-pix2pix ----------------
Copyright (c) 2017, Jun-Yan Zhu and Taesung Park
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

* Redistributions of source code must retain the above copyright notice, this
  list of conditions and the following disclaimer.

* Redistributions in binary form must reproduce the above copyright notice,
  this list of conditions and the following disclaimer in the documentation
  and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

================================================
FILE: README.md
================================================
# maua-stylegan2

This is the repo for my experiments with StyleGAN2. There are many like it, but this one is mine.

It contains the code for [Audio-reactive Latent Interpolations with StyleGAN](https://wavefunk.xyz/assets/audio-reactive-stylegan/paper.pdf) for the NeurIPS 2020 [Workshop on Machine Learning for Creativity and Design](https://neurips2020creativity.github.io/).

The original base is [Kim Seonghyeon's excellent implementation](https://github.com/rosinality/stylegan2-pytorch), but I've gathered code from multiple different repositories or other places online and hacked/grafted it all together. License information for the code should all be in the LICENSE folder, but if you find anything missing or incorrect please let me know and I'll fix it immediately. Tread carefully when trying to distribute any code from this repo, it's meant for research and demonstration.

The files/folders of interest and their purpose are:

| File/Folder | Description
| :--- | :----------
| generate_audiovisual.py | used to generate audio-reactive interpolations
| audioreactive/ | contains the main functions needed for audioreactiveness + examples demonstrating how they can be used
| render.py | renders interpolations using ffmpeg
| select_latents.py | GUI for selecting latents, left click to add to top set, right click to add to bottom
| models/ | StyleGAN networks
| workspace/ | place to store intermediate results, latents, or inputs, etc.
| output/ | default generated output folder
| train.py | code for training models

The rest of the code is experimental, probably broken, and unsupported.

## Installation

```bash
git clone https://github.com/JCBrouwer/maua-stylegan2
cd maua-stylegan2
pip install -r requirements.txt
```

Alternatively, check out this [Colab Notebook](https://colab.research.google.com/drive/1Ig1EXfmBC01qik11Q32P0ZffFtNipiBR)

## Generating audio-reactive interpolations

The simplest way to get started is to try either (in shell):
```bash
python generate_audiovisual.py --ckpt "/path/to/model.pt" --audio_file "/path/to/audio.wav"
```
or (in e.g. a jupyter notebook):
```python
from generate_audiovisual import generate
generate("/path/to/model.pt", "/path/to/audio.wav")
```

This will use the default audio-reactive settings (which aren't great).

To customize the generated interpolation, more functions can be defined to generate latents, noise, network bends, model rewrites, and truncation.

```python
import audioreactive as ar
from generate_audiovisual import generate

def initialize(args):
    args.onsets = ar.onsets(args.audio, args.sr, ...)
    args.chroma = ar.chroma(args.audio, args.sr, ...)
    return args

def get_latents(selection, args):
    latents = ar.chroma_weight_latents(args.chroma, selection)
    return latents

def get_noise(height, width, scale, num_scales, args):
    noise = ar.perlin_noise(...)
    noise *= 1 + args.onsets
    return noise

generate(ckpt="/path/to/model.pt", audio_file="/path/to/audio.wav", initialize=initialize, get_latents=get_latents, get_noise=get_noise)
```

When running from command line, the `generate()` call at the end can be left out and the interpolation can be generated with:

```bash
python generate_audiovisual.py --ckpt "/path/to/model.pt" --audio_file "/path/to/audio.wav" --audioreactive_file "/path/to/the/code_above.py"
```

This lets you change arguments on the command line rather than having to add them to the `generate()` call in you python file (use whatever you prefer).

Within these functions, you can execute any python code to make the inputs to the network react to the music. There are a number of useful functions provided in `audioreactive/` (imported above as `ar`).

Examples showing how to use the library and demonstrating some of the techniques discussed in the paper can be found in `audioreactive/examples/`. A playlist with example results can be found [here](https://www.youtube.com/watch?v=2LxHRGppdpA&list=PLkain1QGMwiWndQwr3U4shvNpoFC21E3a).

One important thing to note is that the outputs of the functions must adhere strictly to the expected formats. 

Each of the functions is called with all of the arguments from the command line (or `generate()`) in the `args` variable. On top of the arguments, `args` also contains:
- audio: raw audio signal
- sr: sampling rate of audio
- n_frames: total number of interpolation frames
- duration: length of audio in seconds

```python
def initialize(args):
    # intialize values used in multiple of the following functions here
    # e.g. onsets, chroma, RMS, segmentations, bpms, etc.
    # this is useful to prevent duplicate computations (get_noise is called for each noise size)
    # remember to store them back in args
    ...
    return args

def get_latents(selection, args):
    # selection holds some latent vectors (generated randomly or from a file)
    # generate an audioreactive latent tensor of shape [n_frames, layers, latent_dim]
    ...
    return latents

def get_noise(height, width, scale, num_scales, args):
    # height and width are the spatial dimensions of the current noise layer
    # scale is the index and num_scales the total number of noise layers
    # generate an audioreactive noise tensor of shape [n_frames, 1, height, width]
    ...
    return noise

def get_bends(args):
    # generate a list of dictionaries specifying network bends
    # these must follow one of two forms:
    #
    # either: {
    #     "layer": layer index to apply bend to,
    #     "transform": torch.nn.Module that applies the transformation,
    # }
    # or: {
    #     "layer": layer index to apply bend to,
    #     "modulation": time dependent modulation of the transformation, shape=(n_frames, ...), 
    #     "transform": function that takes a batch of modulation and returns a torch.nn.Module
    #                  that applies the transformation (given the modulation batch),
    # }
    # (The second one is technical debt in a nutshell. It's a workaround to get kornia transforms
    #  to play nicely. You're probably better off using the first option with a th.nn.Module that
    #  has its modulation as an attribute and keeps count of which frame it's rendering internally).
    ...
    return bends

def get_rewrites(args):
    # generate a dictionary specifying model rewrites
    # each key value pair should follow:
    #       param_name -> [transform, modulation]
    # where: param_name is the fully-qualified parameter name (see generator.named_children())
    #        transform & modulation follow the form of the second network bending dict option above
    ...
    return rewrites

def get_truncation(args):
    # generate a sequence of truncation values of shape (n_frames,)
    ...
    return truncation
```

The arguments to `generate_audiovisual.py` are as follows. The first two are required, and the remaining are optional.
```bash
generate_audiovisual.py
  --ckpt CKPT                              # path to model checkpoint
  --audio_file AUDIO_FILE                  # path to audio file to react to
  --audioreactive_file AUDIOREACTIVE_FILE  # file with audio-reactive functions defined (as above)
  --output_dir OUTPUT_DIR                  # path to output dir
  --offset OFFSET                          # starting time in audio in seconds (defaults to 0)
  --duration DURATION                      # duration of interpolation to generate in seconds (leave empty for length of audiofile)
  --latent_file LATENT_FILE                # path to latents saved as numpy array
  --shuffle_latents                        # whether to shuffle the supplied latents or not
  --out_size OUT_SIZE                      # ouput video size: [512, 1024, or 1920]
  --fps FPS                                # output video framerate
  --batch BATCH                            # batch size to render with
  --truncation TRUNCATION                  # truncation to render with (leave empty if get_truncations() is in --audioreactive_file)
  --randomize_noise                        # whether to randomize noise
  --dataparallel                           # whether to use data parallel rendering
  --stylegan1                              # if the model checkpoint is StyleGAN1
  --G_res G_RES                            # training resolution of the generator
  --base_res_factor BASE_RES_FACTOR        # factor to increase generator noise maps by (useful when e.g. doubling 512px net to 1024px)
  --noconst                                # whether the generator was trained without a constant input layer
  --latent_dim LATENT_DIM                  # latent vector size of the generator
  --n_mlp N_MLP                            # number of mapping network layers
  --channel_multiplier CHANNEL_MULTIPLIER  # generator's channel scaling multiplier
```

Alternatively, `generate()` can be called directly from python. It takes the same arguments as generate_audiovisual.py except instead of supplying an audioreactive_file, the functions should be supplied directly (i.e. initialize, get_latents, get_noise, get_bends, get_rewrites, and get_truncation as arguments).

Model checkpoints can be converted from tensorflow .pkl's with [Kim Seonghyeon's script](https://github.com/rosinality/stylegan2-pytorch/blob/master/convert_weight.py) (the one in this repo is broken). Both StyleGAN2 and StyleGAN2-ADA tensorflow checkpoints should work once converted. A good place to find models is [this repo](https://github.com/justinpinkney/awesome-pretrained-stylegan2).

There is minimal support for rendering with StyleGAN1 checkpoints as well, although only with latent and noise (no network bending or model rewriting).

## Citation

If you use the techniques introduced in the paper or the code in this repository for your research, please cite the paper:
```
@InProceedings{Brouwer_2020_NeurIPS_Workshops},
    author = {Brouwer, Hans},
    title = {Audio-reactive Latent Interpolations with StyleGAN},
    booktitle = {Proceedings of the 4th Workshop on Machine Learning for Creativity and Design at NeurIPS 2020},
    month = {December},
    year = {2020},
    url={https://jcbrouwer.github.io/assets/audio-reactive-stylegan/paper.pdf}
}
```


================================================
FILE: accelerate/accelerate_inception.py
================================================
import os
import gc
import wandb
import argparse
import torch as th
from tqdm import tqdm
from torch.utils import data
import torch.nn.functional as F
from inception_vae import InceptionVAE
from dataset import MultiResolutionDataset
from torchvision import transforms, utils, models


def info(x):
    print(x.shape, x.detach().cpu().min(), x.detach().cpu().mean(), x.detach().cpu().max())


def sample_data(loader):
    while True:
        for batch in loader:
            yield batch


class VGG19(th.nn.Module):
    """
    Adapted from https://github.com/NVIDIA/pix2pixHD
    See LICENSE-VGG
    """

    def __init__(self, requires_grad=False):
        super(VGG19, self).__init__()
        vgg_pretrained_features = models.vgg19(pretrained=True).features
        self.slice1 = th.nn.Sequential()
        self.slice2 = th.nn.Sequential()
        self.slice3 = th.nn.Sequential()
        self.slice4 = th.nn.Sequential()
        self.slice5 = th.nn.Sequential()
        for x in range(2):
            self.slice1.add_module(str(x), vgg_pretrained_features[x])
        for x in range(2, 7):
            self.slice2.add_module(str(x), vgg_pretrained_features[x])
        for x in range(7, 12):
            self.slice3.add_module(str(x), vgg_pretrained_features[x])
        for x in range(12, 21):
            self.slice4.add_module(str(x), vgg_pretrained_features[x])
        for x in range(21, 30):
            self.slice5.add_module(str(x), vgg_pretrained_features[x])
        if not requires_grad:
            for param in self.parameters():
                param.requires_grad = False

    def forward(self, X):
        h_relu1 = self.slice1(X)
        h_relu2 = self.slice2(h_relu1)
        h_relu3 = self.slice3(h_relu2)
        h_relu4 = self.slice4(h_relu3)
        h_relu5 = self.slice5(h_relu4)
        out = [h_relu1, h_relu2, h_relu3, h_relu4, h_relu5]
        return out


class VGGLoss(th.nn.Module):
    """
    Adapted from https://github.com/NVIDIA/pix2pixHD
    See LICENSE-VGG
    """

    def __init__(self):
        super(VGGLoss, self).__init__()
        self.vgg = VGG19().cuda()
        self.criterion = th.nn.L1Loss()
        self.weights = [1.0 / 32, 1.0 / 16, 1.0 / 8, 1.0 / 4, 1.0]

    def forward(self, x, y):
        x_vgg, y_vgg = self.vgg(x), self.vgg(y)
        loss = 0
        for i in range(len(x_vgg)):
            loss += self.weights[i] * self.criterion(x_vgg[i], y_vgg[i].detach())
        return loss


def train(latent_dim, num_repeats, learning_rate, lambda_vgg, lambda_mse):
    print(
        f"latent_dim={latent_dim:.4f}",
        f"num_repeats={num_repeats:.4f}",
        f"learning_rate={learning_rate:.4f}",
        f"lambda_vgg={lambda_vgg:.4f}",
        f"lambda_mse={lambda_mse:.4f}",
    )

    transform = transforms.Compose(
        [
            transforms.Resize(128),
            transforms.RandomHorizontalFlip(p=0.5),
            transforms.ToTensor(),
            # transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True),
        ]
    )
    batch_size = 72
    data_path = "/home/hans/trainsets/cyphis"
    name = os.path.splitext(os.path.basename(data_path))[0]
    dataset = MultiResolutionDataset(data_path, transform, 256)
    dataloader = data.DataLoader(
        dataset, batch_size=batch_size, sampler=data.RandomSampler(dataset), num_workers=12, drop_last=True,
    )
    loader = sample_data(dataloader)
    sample_imgs = next(loader)[:24]
    wandb.log({"Real Images": [wandb.Image(utils.make_grid(sample_imgs, nrow=6, normalize=True, range=(0, 1)))]})

    vae, vae_optim = None, None
    vae = InceptionVAE(latent_dim=latent_dim, repeat_per_block=num_repeats).to(device)
    vae_optim = th.optim.Adam(vae.parameters(), lr=learning_rate)

    vgg = VGGLoss()

    # sample_z = th.randn(size=(24, 512))

    scores = []
    num_iters = 100_000
    pbar = tqdm(range(num_iters), smoothing=0.1)
    for i in pbar:
        vae.train()

        real = next(loader).to(device)

        fake, mu, log_var = vae(real)

        bce = F.binary_cross_entropy(fake, real, size_average=False)
        kld = -0.5 * th.sum(1 + log_var - mu.pow(2) - log_var.exp())
        vgg_loss = vgg(fake, real)
        mse_loss = th.sqrt((fake - real).pow(2).mean())

        loss = bce + kld + lambda_vgg * vgg_loss + lambda_mse * mse_loss

        loss_dict = {
            "Total": loss,
            "BCE": bce,
            "Kullback Leibler Divergence": kld,
            "MSE": mse_loss,
            "VGG": vgg_loss,
        }

        vae.zero_grad()
        loss.backward()
        vae_optim.step()

        wandb.log(loss_dict)

        with th.no_grad():
            if i % int(num_iters / 100) == 0 or i + 1 == num_iters:
                vae.eval()

                sample, _, _ = vae(sample_imgs.to(device))
                grid = utils.make_grid(sample, nrow=6, normalize=True, range=(0, 1))
                del sample
                wandb.log({"Reconstructed Images VAE": [wandb.Image(grid, caption=f"Step {i}")]})

                sample = vae.sampling()
                grid = utils.make_grid(sample, nrow=6, normalize=True, range=(0, 1))
                del sample
                wandb.log({"Generated Images VAE": [wandb.Image(grid, caption=f"Step {i}")]})

                gc.collect()
                th.cuda.empty_cache()

                th.save(
                    {"vae": vae.state_dict(), "vae_optim": vae_optim.state_dict()},
                    f"/home/hans/modelzoo/maua-sg2/vae-{name}-{wandb.run.dir.split('/')[-1].split('-')[-1]}.pt",
                )

        if th.isnan(loss).any() or th.isinf(loss).any():
            print("NaN losses, exiting...")
            print(
                {
                    "Total": loss,
                    "\nBCE": bce,
                    "\nKullback Leibler Divergence": kld,
                    "\nMSE": mse_loss,
                    "\nVGG": vgg_loss,
                }
            )
            wandb.log({"Total": 27000})
            return


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--latent_dim", type=float, default=512)
    parser.add_argument("--num_repeats", type=float, default=1)
    parser.add_argument("--learning_rate", type=float, default=0.005)
    parser.add_argument("--lambda_vgg", type=float, default=1.0)
    parser.add_argument("--lambda_mse", type=float, default=1.0)
    args = parser.parse_args()

    device = "cuda"
    th.backends.cudnn.benchmark = True

    wandb.init(project=f"maua-stylegan")

    train(
        args.latent_dim, args.num_repeats, args.learning_rate, args.lambda_vgg, args.lambda_mse,
    )



================================================
FILE: accelerate/accelerate_logcosh.py
================================================
import os
import gc
import wandb
import argparse
import validation
import torch as th
from tqdm import tqdm
from torch.utils import data
from autoencoder import LogCoshVAE
from dataset import MultiResolutionDataset
from torchvision import transforms, utils, models


def data_sampler(dataset, shuffle, distributed):
    if distributed:
        return data.distributed.DistributedSampler(dataset, shuffle=shuffle)
    if shuffle:
        return data.RandomSampler(dataset)
    else:
        return data.SequentialSampler(dataset)


def sample_data(loader):
    while True:
        for batch in loader:
            yield batch


class VGG19(th.nn.Module):
    """
    Adapted from https://github.com/NVIDIA/pix2pixHD
    See LICENSE-VGG
    """

    def __init__(self, requires_grad=False):
        super(VGG19, self).__init__()
        vgg_pretrained_features = models.vgg19(pretrained=True).features
        self.slice1 = th.nn.Sequential()
        self.slice2 = th.nn.Sequential()
        self.slice3 = th.nn.Sequential()
        self.slice4 = th.nn.Sequential()
        self.slice5 = th.nn.Sequential()
        for x in range(2):
            self.slice1.add_module(str(x), vgg_pretrained_features[x])
        for x in range(2, 7):
            self.slice2.add_module(str(x), vgg_pretrained_features[x])
        for x in range(7, 12):
            self.slice3.add_module(str(x), vgg_pretrained_features[x])
        for x in range(12, 21):
            self.slice4.add_module(str(x), vgg_pretrained_features[x])
        for x in range(21, 30):
            self.slice5.add_module(str(x), vgg_pretrained_features[x])
        if not requires_grad:
            for param in self.parameters():
                param.requires_grad = False

    def forward(self, X):
        h_relu1 = self.slice1(X)
        h_relu2 = self.slice2(h_relu1)
        h_relu3 = self.slice3(h_relu2)
        h_relu4 = self.slice4(h_relu3)
        h_relu5 = self.slice5(h_relu4)
        out = [h_relu1, h_relu2, h_relu3, h_relu4, h_relu5]
        return out


class VGGLoss(th.nn.Module):
    """
    Adapted from https://github.com/NVIDIA/pix2pixHD
    See LICENSE-VGG
    """

    def __init__(self):
        super(VGGLoss, self).__init__()
        self.vgg = VGG19().cuda()
        self.criterion = th.nn.L1Loss()
        self.weights = [1.0 / 32, 1.0 / 16, 1.0 / 8, 1.0 / 4, 1.0]

    def forward(self, x, y):
        x_vgg, y_vgg = self.vgg(x), self.vgg(y)
        loss = 0
        for i in range(len(x_vgg)):
            loss += self.weights[i] * self.criterion(x_vgg[i], y_vgg[i].detach())
        return loss


device = "cuda"
th.backends.cudnn.benchmark = True

wandb.init(project=f"maua-stylegan")


def train(latent_dim, learning_rate, number_filters, vae_alpha, vae_beta, kl_divergence_weight):
    print(
        f"latent_dim={latent_dim}",
        f"learning_rate={learning_rate}",
        f"number_filters={number_filters}",
        f"vae_alpha={vae_alpha}",
        f"vae_beta={vae_beta}",
        f"kl_divergence_weight={kl_divergence_weight}",
    )

    batch_size = 64
    i = None
    while batch_size >= 1:
        try:
            transform = transforms.Compose(
                [
                    transforms.Resize(128),
                    transforms.RandomHorizontalFlip(p=0.5),
                    transforms.ToTensor(),
                    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True),
                ]
            )
            data_path = "/home/hans/trainsets/cyphis"
            name = os.path.splitext(os.path.basename(data_path))[0]
            dataset = MultiResolutionDataset(data_path, transform, 256)
            dataloader = data.DataLoader(
                dataset,
                batch_size=int(batch_size),
                sampler=data_sampler(dataset, shuffle=True, distributed=False),
                num_workers=12,
                drop_last=True,
            )
            loader = sample_data(dataloader)
            sample_imgs = next(loader)[:24]
            wandb.log(
                {"Real Images": [wandb.Image(utils.make_grid(sample_imgs, nrow=6, normalize=True, range=(-1, 1)))]}
            )

            hidden_dims = [min(int(number_filters) * 2 ** i, latent_dim) for i in range(5)] + [latent_dim]
            vae, vae_optim = None, None
            vae = LogCoshVAE(
                3, latent_dim, hidden_dims=hidden_dims, alpha=vae_alpha, beta=vae_beta, kld_weight=kl_divergence_weight,
            ).to(device)
            vae.train()
            vae_optim = th.optim.Adam(vae.parameters(), lr=learning_rate)

            mse_loss = th.nn.MSELoss()
            vgg = VGGLoss()

            sample_z = th.randn(size=(24, latent_dim))

            scores = []
            num_iters = 100_000
            pbar = range(num_iters)
            pbar = tqdm(pbar, smoothing=0.1)
            for i in pbar:
                vae.train()

                real = next(loader).to(device)
                fake, mu, log_var = vae(real)

                loss_dict = vae.loss(real, fake, mu, log_var)
                vgg_loss = vgg(fake, real)
                loss = loss_dict["Total"] + vgg_loss

                vae.zero_grad()
                loss.backward()
                vae_optim.step()

                wandb.log(
                    {
                        "Total": loss,
                        "VGG": vgg_loss,
                        "Reconstruction": loss_dict["Reconstruction"],
                        "Kullback Leibler Divergence": loss_dict["Kullback Leibler Divergence"],
                    }
                )

                if i % int(num_iters / 1000) == 0 or i + 1 == num_iters:
                    with th.no_grad():
                        vae.eval()

                        sample, _, _ = vae(sample_imgs.to(device))
                        grid = utils.make_grid(sample, nrow=6, normalize=True, range=(-1, 1),)
                        del sample
                        wandb.log({"Reconstructed Images VAE": [wandb.Image(grid, caption=f"Step {i}")]})

                        sample = vae.decode(sample_z.to(device))
                        grid = utils.make_grid(sample, nrow=6, normalize=True, range=(-1, 1),)
                        del sample
                        wandb.log({"Generated Images VAE": [wandb.Image(grid, caption=f"Step {i}")]})

                if i % int(num_iters / 40) == 0 or i + 1 == num_iters:
                    with th.no_grad():
                        fid_dict = validation.vae_fid(vae, int(batch_size), (latent_dim,), 5000, name)
                        wandb.log(fid_dict)
                        mse = mse_loss(fake, real) * 5000
                        score = fid_dict["FID"] + mse + 1000 * vgg_loss
                        wandb.log({"Score": score})
                        pbar.set_description(f"FID: {fid_dict['FID']:.2f} MSE: {mse:.2f} VGG: {1000 * vgg_loss:.2f}")

                    if i >= num_iters / 2:
                        scores.append(score)

                if th.isnan(loss).any() or th.isinf(loss).any():
                    print("NaN losses, exiting...")
                    print(
                        {
                            "Total": loss.detach().cpu().item(),
                            "\nVGG": vgg_loss.detach().cpu().item(),
                            "\nReconstruction": loss_dict["Reconstruction"].detach().cpu().item(),
                            "\nKullback Leibler Divergence": loss_dict["Kullback Leibler Divergence"]
                            .detach()
                            .cpu()
                            .item(),
                        }
                    )
                    wandb.log({"Score": 27000})
                    return

            return

        except RuntimeError as e:
            if "CUDA out of memory" in str(e):
                batch_size = batch_size / 2

                if batch_size < 1:
                    print("This configuration does not fit into memory, exiting...")
                    wandb.log({"Score": 27000})
                    return

                print(f"Out of memory, halving batch size... {batch_size}")
                if vae is not None:
                    del vae
                if vae_optim is not None:
                    del vae_optim
                gc.collect()
                th.cuda.empty_cache()

            else:
                print(e)
                return


parser = argparse.ArgumentParser()
parser.add_argument("--latent_dim", type=int, default=1024)
parser.add_argument("--learning_rate", type=float, default=0.005)
parser.add_argument("--number_filters", type=int, default=64)
parser.add_argument("--vae_alpha", type=float, default=10.0)
parser.add_argument("--vae_beta", type=float, default=1.0)
parser.add_argument("--kl_divergence_weight", type=float, default=1.0)
args = parser.parse_args()

train(
    args.latent_dim, args.learning_rate, args.number_filters, args.vae_alpha, args.vae_beta, args.kl_divergence_weight,
)



================================================
FILE: accelerate/accelerate_segnet.py
================================================
import os
import gc
import wandb
import argparse
import torch as th
from tqdm import tqdm
from torch.utils import data
from autoencoder import ConvSegNet
from dataset import MultiResolutionDataset
from torchvision import transforms, utils, models


def info(x):
    print(x.shape, x.detach().cpu().min(), x.detach().cpu().mean(), x.detach().cpu().max())


def sample_data(loader):
    while True:
        for batch in loader:
            yield batch


class VGG19(th.nn.Module):
    """
    Adapted from https://github.com/NVIDIA/pix2pixHD
    See LICENSE-VGG
    """

    def __init__(self, requires_grad=False):
        super(VGG19, self).__init__()
        vgg_pretrained_features = models.vgg19(pretrained=True).features
        self.slice1 = th.nn.Sequential()
        self.slice2 = th.nn.Sequential()
        self.slice3 = th.nn.Sequential()
        self.slice4 = th.nn.Sequential()
        self.slice5 = th.nn.Sequential()
        for x in range(2):
            self.slice1.add_module(str(x), vgg_pretrained_features[x])
        for x in range(2, 7):
            self.slice2.add_module(str(x), vgg_pretrained_features[x])
        for x in range(7, 12):
            self.slice3.add_module(str(x), vgg_pretrained_features[x])
        for x in range(12, 21):
            self.slice4.add_module(str(x), vgg_pretrained_features[x])
        for x in range(21, 30):
            self.slice5.add_module(str(x), vgg_pretrained_features[x])
        if not requires_grad:
            for param in self.parameters():
                param.requires_grad = False

    def forward(self, X):
        h_relu1 = self.slice1(X)
        h_relu2 = self.slice2(h_relu1)
        h_relu3 = self.slice3(h_relu2)
        h_relu4 = self.slice4(h_relu3)
        h_relu5 = self.slice5(h_relu4)
        out = [h_relu1, h_relu2, h_relu3, h_relu4, h_relu5]
        return out


class VGGLoss(th.nn.Module):
    """
    Adapted from https://github.com/NVIDIA/pix2pixHD
    See LICENSE-VGG
    """

    def __init__(self):
        super(VGGLoss, self).__init__()
        self.vgg = VGG19().cuda()
        self.criterion = th.nn.L1Loss()
        self.weights = [1.0 / 32, 1.0 / 16, 1.0 / 8, 1.0 / 4, 1.0]

    def forward(self, x, y):
        x_vgg, y_vgg = self.vgg(x), self.vgg(y)
        loss = 0
        for i in range(len(x_vgg)):
            loss += self.weights[i] * self.criterion(x_vgg[i], y_vgg[i].detach())
        return loss


def align(x, y, alpha=2):
    return (x - y).norm(p=2, dim=1).pow(alpha).mean()


def uniform(x, t=2):
    return (th.pdist(x.view(x.size(0), -1), p=2).pow(2).mul(-t).exp().mean() + 1e-27).log()


def train(learning_rate, lambda_mse):
    print(
        f"learning_rate={learning_rate:.4f}", f"lambda_mse={lambda_mse:.4f}",
    )

    transform = transforms.Compose(
        [
            transforms.Resize(128),
            transforms.RandomHorizontalFlip(p=0.5),
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True),
        ]
    )
    batch_size = 72
    data_path = "/home/hans/trainsets/cyphis"
    name = os.path.splitext(os.path.basename(data_path))[0]
    dataset = MultiResolutionDataset(data_path, transform, 256)
    dataloader = data.DataLoader(
        dataset, batch_size=batch_size, sampler=data.RandomSampler(dataset), num_workers=12, drop_last=True,
    )
    loader = sample_data(dataloader)
    sample_imgs = next(loader)[:24]
    wandb.log({"Real Images": [wandb.Image(utils.make_grid(sample_imgs, nrow=6, normalize=True, range=(-1, 1)))]})

    vae, vae_optim = None, None
    vae = ConvSegNet().to(device)
    vae_optim = th.optim.Adam(vae.parameters(), lr=learning_rate)

    vgg = VGGLoss()

    sample_z = th.randn(size=(24, 512, 16, 16))
    sample_z /= sample_z.abs().max()

    scores = []
    num_iters = 100_000
    pbar = tqdm(range(num_iters), smoothing=0.1)
    for i in pbar:
        vae.train()

        real = next(loader).to(device)

        z = vae.encode(real)
        fake = vae.decode(z)

        vgg_loss = vgg(fake, real)

        mse_loss = th.sqrt((fake - real).pow(2).mean())

        # diff = fake - real
        # recons_loss = recons_alpha * diff + th.log(1.0 + th.exp(-2 * recons_alpha * diff)) - th.log(th.tensor(2.0))
        # recons_loss = (1.0 / recons_alpha) * recons_loss.mean()
        # recons_loss = recons_loss if not th.isinf(recons_loss).any() else 0

        # x, y = z.chunk(2)
        # align_loss = align(x, y, alpha=align_alpha)
        # unif_loss = -(uniform(x, t=unif_t) + uniform(y, t=unif_t)) / 2.0

        loss = (
            vgg_loss
            + lambda_mse * mse_loss
            # + lambda_recons * recons_loss
            # + lambda_align * align_loss
            # + lambda_unif * unif_loss
        )
        # print(vgg_loss.detach().cpu().item())
        # print(lambda_mse * mse_loss.detach().cpu().item())
        # # print(lambda_recons * recons_loss.detach().cpu().item())
        # print(lambda_align * align_loss.detach().cpu().item())
        # print(lambda_unif * unif_loss.detach().cpu().item())

        loss_dict = {
            "Total": loss,
            "MSE": mse_loss,
            "VGG": vgg_loss,
            # "Reconstruction": recons_loss,
            # "Alignment": align_loss,
            # "Uniformity": unif_loss,
        }

        vae.zero_grad()
        loss.backward()
        vae_optim.step()

        wandb.log(loss_dict)
        # pbar.set_description(" ".join())

        with th.no_grad():
            if i % int(num_iters / 100) == 0 or i + 1 == num_iters:
                vae.eval()

                sample = vae(sample_imgs.to(device))
                grid = utils.make_grid(sample, nrow=6, normalize=True, range=(-1, 1))
                del sample
                wandb.log({"Reconstructed Images VAE": [wandb.Image(grid, caption=f"Step {i}")]})

                sample = vae.decode(sample_z.to(device))
                grid = utils.make_grid(sample, nrow=6, normalize=True, range=(-1, 1))
                del sample
                wandb.log({"Generated Images VAE": [wandb.Image(grid, caption=f"Step {i}")]})

                gc.collect()
                th.cuda.empty_cache()

                th.save(
                    {"vae": vae.state_dict(), "vae_optim": vae_optim.state_dict()},
                    f"/home/hans/modelzoo/maua-sg2/vae-{name}-{wandb.run.dir.split('/')[-1].split('-')[-1]}.pt",
                )

        if th.isnan(loss).any():
            print("NaN losses, exiting...")
            wandb.log({"Total": 27000})
            return


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--learning_rate", type=float, default=0.005)
    parser.add_argument("--lambda_mse", type=float, default=1.0)
    # parser.add_argument("--lambda_recons", type=float, default=0.0)
    # parser.add_argument("--recons_alpha", type=float, default=5.0)
    # parser.add_argument("--lambda_align", type=float, default=1.0)
    # parser.add_argument("--align_alpha", type=float, default=2.0)
    # parser.add_argument("--lambda_unif", type=float, default=1.0)
    # parser.add_argument("--unif_t", type=float, default=0.001)
    args = parser.parse_args()

    device = "cuda"
    th.backends.cudnn.benchmark = True

    wandb.init(project=f"maua-stylegan")

    train(
        args.learning_rate,
        args.lambda_mse,
        # args.lambda_recons,
        # args.recons_alpha,
        # args.lambda_align,
        # args.align_alpha,
        # args.lambda_unif,
        # args.unif_t,
    )



================================================
FILE: audioreactive/__init__.py
================================================
from .bend import *
from .examples import *
from .latent import *
from .signal import *
from .util import *


================================================
FILE: audioreactive/bend.py
================================================
import math

import kornia.augmentation as kA
import kornia.geometry.transform as kT
import torch as th

# ====================================================================================
# ================================= network bending ==================================
# ====================================================================================


class NetworkBend(th.nn.Module):
    """Base network bending class

    Args:
        sequential_fn (function): Function that takes a batch of modulation and creates th.nn.Sequential
        modulation (th.tensor): Modulation batch
    """

    def __init__(self, sequential_fn, modulation):
        super(NetworkBend, self).__init__()
        self.sequential = sequential_fn(modulation)

    def forward(self, x):
        return self.sequential(x)


class AddNoise(th.nn.Module):
    """Adds static noise to output

    Args:
        noise (th.tensor): Noise to be added
    """

    def __init__(self, noise):
        super(AddNoise, self).__init__()
        self.noise = noise

    def forward(self, x):
        return x + self.noise.to(x.device)


class Print(th.nn.Module):
    """Prints intermediate feature statistics (useful for debugging complicated network bends)."""

    def forward(self, x):
        print(x.shape, [x.min().item(), x.mean().item(), x.max().item()], th.std(x).item())
        return x


class Translate(NetworkBend):
    """Creates horizontal translating effect where repeated linear interpolations from 0 to 1 (saw tooth wave) creates seamless scrolling effect.

    Args:
        modulation (th.tensor): [0.0-1.0]. Batch of modulation
        h (int): Height of intermediate features that the network bend is applied to
        w (int): Width of intermediate features that the network bend is applied to
        noise (int): Noise to be added (must be 5 * width wide)
    """

    def __init__(self, modulation, h, w, noise):
        sequential_fn = lambda b: th.nn.Sequential(
            th.nn.ReflectionPad2d((int(w / 2), int(w / 2), 0, 0)),
            th.nn.ReflectionPad2d((w, w, 0, 0)),
            th.nn.ReflectionPad2d((w, 0, 0, 0)),
            AddNoise(noise),
            kT.Translate(b),
            kA.CenterCrop((h, w)),
        )
        super(Translate, self).__init__(sequential_fn, modulation)


class Zoom(NetworkBend):
    """Creates zooming effect.

    Args:
        modulation (th.tensor): [0.0-1.0]. Batch of modulation
        h (int): height of intermediate features that the network bend is applied to
        w (int): width of intermediate features that the network bend is applied to
    """

    def __init__(self, modulation, h, w):
        padding = int(max(h, w)) - 1
        sequential_fn = lambda b: th.nn.Sequential(th.nn.ReflectionPad2d(padding), kT.Scale(b), kA.CenterCrop((h, w)))
        super(Zoom, self).__init__(sequential_fn, modulation)


class Rotate(NetworkBend):
    """Creates rotation effect.

    Args:
        modulation (th.tensor): [0.0-1.0]. Batch of modulation
        h (int): height of intermediate features that the network bend is applied to
        w (int): width of intermediate features that the network bend is applied to
    """

    def __init__(self, modulation, h, w):
        # worst case rotation brings sqrt(2) * max_side_length out-of-frame pixels into frame
        # padding should cover that exactly
        padding = int(max(h, w) * (1 - math.sqrt(2) / 2))
        sequential_fn = lambda b: th.nn.Sequential(th.nn.ReflectionPad2d(padding), kT.Rotate(b), kA.CenterCrop((h, w)))
        super(Rotate, self).__init__(sequential_fn, modulation)


================================================
FILE: audioreactive/examples/__init__.py
================================================
from . import *


================================================
FILE: audioreactive/examples/default.py
================================================
import torch as th

import audioreactive as ar


def initialize(args):
    args.lo_onsets = ar.onsets(args.audio, args.sr, args.n_frames, fmax=150, smooth=5, clip=97, power=2)
    args.hi_onsets = ar.onsets(args.audio, args.sr, args.n_frames, fmin=500, smooth=5, clip=99, power=2)
    return args


def get_latents(selection, args):
    chroma = ar.chroma(args.audio, args.sr, args.n_frames)
    chroma_latents = ar.chroma_weight_latents(chroma, selection)
    latents = ar.gaussian_filter(chroma_latents, 4)

    lo_onsets = args.lo_onsets[:, None, None]
    hi_onsets = args.hi_onsets[:, None, None]

    latents = hi_onsets * selection[[-4]] + (1 - hi_onsets) * latents
    latents = lo_onsets * selection[[-7]] + (1 - lo_onsets) * latents

    latents = ar.gaussian_filter(latents, 2, causal=0.2)

    return latents


def get_noise(height, width, scale, num_scales, args):
    if width > 256:
        return None

    lo_onsets = args.lo_onsets[:, None, None, None].cuda()
    hi_onsets = args.hi_onsets[:, None, None, None].cuda()

    noise_noisy = ar.gaussian_filter(th.randn((args.n_frames, 1, height, width), device="cuda"), 5)
    noise = ar.gaussian_filter(th.randn((args.n_frames, 1, height, width), device="cuda"), 128)

    if width < 128:
        noise = lo_onsets * noise_noisy + (1 - lo_onsets) * noise
    if width > 32:
        noise = hi_onsets * noise_noisy + (1 - hi_onsets) * noise

    noise /= noise.std() * 2.5

    return noise.cpu()


================================================
FILE: audioreactive/examples/kelp.py
================================================
"""
This file shows an example of a loop based interpolation
Here sections are identified with laplacian segmentation and looping latents are generated for each section
The noise is looping perlin noise
Long term section analysis is done with the RMS to interpolate between latent sequences for the intro/outro and drop
"""


import librosa as rosa
import torch as th

import audioreactive as ar

OVERRIDE = dict(audio_file="audioreactive/examples/Wavefunk - Dwelling in the Kelp.mp3", out_size=1920)
BPM = 130


def initialize(args):
    # RMS can be used to distinguish between the drop sections and intro/outros
    rms = ar.rms(args.audio, args.sr, args.n_frames, smooth=10, clip=60, power=1)
    rms = ar.expand(rms, threshold=0.8, ratio=10)
    rms = ar.gaussian_filter(rms, 4)
    rms = ar.normalize(rms)
    args.rms = rms

    # cheating a little here, this my song so I have the multitracks
    # this is much easier than fiddling with onsets until you have envelopes that dance nicely to the drums
    audio, sr = rosa.load("workspace/kelpkick.wav", offset=args.offset, duration=args.duration)
    args.kick_onsets = ar.onsets(audio, sr, args.n_frames, margin=1, smooth=4)
    audio, sr = rosa.load("workspace/kelpsnare.wav", offset=args.offset, duration=args.duration)
    args.snare_onsets = ar.onsets(audio, sr, args.n_frames, margin=1, smooth=4)

    ar.plot_signals([args.rms, args.kick_onsets, args.snare_onsets])

    return args


def get_latents(selection, args):
    # expand envelopes to latent shape
    rms = args.rms[:, None, None]
    low_onsets = args.kick_onsets[:, None, None]
    high_onsets = args.snare_onsets[:, None, None]

    # get timestamps and labels with laplacian segmentation
    # k is the number of labels the algorithm may use
    # try multiple values with plot=True to see which value correlates best with the sections of the song
    timestamps, labels = ar.laplacian_segmentation(args.audio, args.sr, k=7)

    # a second set of latents for the drop section, the 'selection' variable is the other set for the intro
    drop_selection = ar.load_latents("workspace/cyphept_kelp_drop_latents.npy")
    color_layer = 9

    latents = []
    for (start, stop), l in zip(zip(timestamps, timestamps[1:]), labels):
        start_frame = int(round(start / args.duration * args.n_frames))
        stop_frame = int(round(stop / args.duration * args.n_frames))
        section_frames = stop_frame - start_frame
        section_bars = (stop - start) * (BPM / 60) / 4

        # get portion of latent selection (wrapping around to start)
        latent_selection_slice = ar.wrapping_slice(selection, l, 4)
        # spline interpolation loops through selection slice
        latent_section = ar.spline_loops(latent_selection_slice, n_frames=section_frames, n_loops=section_bars / 4)
        # set the color with laplacian segmentation label, (1 latent repeated for entire section in upper layers)
        latent_section[:, color_layer:] = th.cat([selection[[l], color_layer:]] * section_frames)

        # same as above but for the drop latents (with faster loops)
        drop_selection_slice = ar.wrapping_slice(drop_selection, l, 4)
        drop_section = ar.spline_loops(drop_selection_slice, n_frames=section_frames, n_loops=section_bars / 2)
        drop_section[:, color_layer:] = th.cat([drop_selection[[l], color_layer:]] * section_frames)

        # merged based on RMS (drop section or not)
        latents.append((1 - rms[start_frame:stop_frame]) * latent_section + rms[start_frame:stop_frame] * drop_section)

    # concatenate latents to correct length & smooth over the junctions
    len_latents = sum([len(l) for l in latents])
    if len_latents != args.n_frames:
        latents.append(th.cat([latents[-1][[-1]]] * (args.n_frames - len_latents)))
    latents = th.cat(latents).float()
    latents = ar.gaussian_filter(latents, 3)

    # use onsets to modulate towards latents
    latents = 0.666 * low_onsets * selection[[2]] + (1 - 0.666 * low_onsets) * latents
    latents = 0.666 * high_onsets * selection[[1]] + (1 - 0.666 * high_onsets) * latents

    latents = ar.gaussian_filter(latents, 1, causal=0.2)
    return latents


def get_noise(height, width, scale, num_scales, args):
    if width > 512:  # larger sizes don't fit in VRAM, just use default or randomize
        return

    num_bars = int(round(args.duration * (BPM / 60) / 4))
    frames_per_loop = int(args.n_frames / num_bars * 2)  # loop every 2 bars

    def perlin_pls(resolution):
        perlin = ar.perlin_noise(shape=(frames_per_loop, height, width), res=resolution)[:, None, ...].cpu()
        perlin = th.cat([perlin] * int(num_bars / 2))  # concatenate multiple copies for looping
        if args.n_frames - len(perlin) > 0:
            perlin = th.cat([perlin, th.cat([perlin[[-1]]] * (args.n_frames - len(perlin)))])  # fix up rounding errors
        return perlin

    smooth = perlin_pls(resolution=(1, 1, 1))  # (time res, x res, y res)
    noise = perlin_pls(resolution=(8, 4, 4))  # higher resolution => higher frequency noise => more movement in video

    rms = args.rms[:, None, None, None]
    noise = rms * noise + (1 - rms) * smooth  # blend between noises based on drop (high rms) or not

    return noise


def get_bends(args):
    # repeat the intermediate features outwards on both sides (2:1 aspect ratio)
    # + add some noise to give the whole thing a little variation (disguises the repetition)
    transform = th.nn.Sequential(
        th.nn.ReplicationPad2d((2, 2, 0, 0)), ar.AddNoise(0.025 * th.randn(size=(1, 1, 4, 8), device="cuda")),
    )
    bends = [{"layer": 0, "transform": transform}]

    return bends


================================================
FILE: audioreactive/examples/tauceti.py
================================================
"""
This file shows an example of network bending
The latents and noise are similar to temper.py (although without spatial noise controls)
The latents cycle through different colors for different sections of the drop
During the drop, a translation is applied which makes the video seem to scroll endlessly
"""

from functools import partial

import numpy as np
import torch as th

import audioreactive as ar

OVERRIDE = dict(
    audio_file="audioreactive/examples/Wavefunk - Tau Ceti Alpha.mp3",
    out_size=1920,  # get bends assumes 1920x1080 output size
    dataparallel=False,  # makes use of a kornia transform during network bending => not compatible with dataparallel
    fps=30,  # 5591 magic number below is based on number of frames in output video with fps of 30
)


def initialize(args):
    args.low_onsets = ar.onsets(args.audio, args.sr, args.n_frames, fmax=150, smooth=5, clip=97, power=2)
    args.high_onsets = ar.onsets(args.audio, args.sr, args.n_frames, fmin=500, smooth=5, clip=99, power=2)
    return args


def get_latents(selection, args):
    chroma = ar.chroma(args.audio, args.sr, args.n_frames)
    chroma_latents = ar.chroma_weight_latents(chroma, selection[:12])  # shape [n_frames, 18, 512]
    latents = ar.gaussian_filter(chroma_latents, 5)

    lo_onsets = args.low_onsets[:, None, None]  # expand to same shape as latents [n_frames, 1, 1]
    hi_onsets = args.high_onsets[:, None, None]

    latents = hi_onsets * selection[[-4]] + (1 - hi_onsets) * latents
    latents = lo_onsets * selection[[-7]] + (1 - lo_onsets) * latents

    latents = ar.gaussian_filter(latents, 5, causal=0)

    # cheating a little, you could probably do this with laplacian segmentation, but is it worth the effort?
    drop_start = int(5591 * (45 / args.duration))
    drop_end = int(5591 * (135 / args.duration))

    # selection of latents with different colors (chosen with select_latents.py)
    color_latent_selection = th.from_numpy(np.load("workspace/cyphept-multicolor-latents.npy"))

    # build sequence of latents for just the upper layers
    color_layer = 9
    color_latents = [latents[:drop_start, color_layer:]]

    # for 4 different sections in the drop, use a different color latent
    drop_length = drop_end - drop_start
    section_length = int(drop_length / 4)
    for i, section_start in enumerate(range(0, drop_length, section_length)):
        if i > 3:
            break
        color_latents.append(th.cat([color_latent_selection[[i], color_layer:]] * section_length))

    # ensure color sequence is correct length and concatenate
    if drop_length - 4 * section_length != 0:
        color_latents.append(th.cat([color_latent_selection[[i], color_layer:]] * (drop_length - 4 * section_length)))
    color_latents.append(latents[drop_end:, color_layer:])
    color_latents = th.cat(color_latents, axis=0)

    color_latents = ar.gaussian_filter(color_latents, 5)

    # set upper layers of latent sequence to the colored sequence
    latents[:, color_layer:] = color_latents

    return latents


def get_noise(height, width, scale, num_scales, args):
    if width > 256:
        return None

    lo_onsets = 1.25 * args.low_onsets[:, None, None, None].cuda()
    hi_onsets = 1.25 * args.high_onsets[:, None, None, None].cuda()

    noise_noisy = ar.gaussian_filter(th.randn((args.n_frames, 1, height, width), device="cuda"), 5)

    noise = ar.gaussian_filter(th.randn((args.n_frames, 1, height, width), device="cuda"), 128)
    if width > 8:
        noise = lo_onsets * noise_noisy + (1 - lo_onsets) * noise
        noise = hi_onsets * noise_noisy + (1 - hi_onsets) * noise

    noise /= noise.std() * 2.5

    return noise.cpu()


def get_bends(args):
    # repeat the intermediate features outwards on both sides (2:1 aspect ratio)
    # + add some noise to give the whole thing a little variation (disguises the repetition)
    transform = th.nn.Sequential(
        th.nn.ReplicationPad2d((2, 2, 0, 0)), ar.AddNoise(0.025 * th.randn(size=(1, 1, 4, 8), device="cuda")),
    )
    bends = [{"layer": 0, "transform": transform}]

    # during the drop, create scrolling effect
    drop_start = int(5591 * (45 / args.duration))
    drop_end = int(5591 * (135 / args.duration))

    # calculate length of loops, number of loops, and remainder at end of drop
    scroll_loop_length = int(6 * args.fps)
    scroll_loop_num = int((drop_end - drop_start) / scroll_loop_length)
    scroll_trunc = (drop_end - drop_start) - scroll_loop_num * scroll_loop_length

    # apply network bending to 4th layer in StyleGAN
    # lower layer network bends have more fluid outcomes
    tl = 4
    h = 2 ** tl
    w = 2 * h

    # create values between 0 and 1 corresponding to fraction of scroll from left to right completed
    # all 0s during intro
    intro_tl8 = np.zeros(drop_start)
    # repeating linear interpolation from 0 to 1 during drop
    loops_tl8 = np.concatenate([np.linspace(0, w, scroll_loop_length)] * scroll_loop_num)
    # truncated interp
    last_loop_tl8 = np.linspace(0, w, scroll_loop_length)[:scroll_trunc]
    # static at final truncated value during outro
    outro_tl8 = np.ones(args.n_frames - drop_end) * np.linspace(0, w, scroll_loop_length)[scroll_trunc + 1]

    # create 2D array of translations in x and y directions
    x_tl8 = np.concatenate([intro_tl8, loops_tl8, last_loop_tl8, outro_tl8])
    y_tl8 = np.zeros(args.n_frames)
    translation = (th.tensor([x_tl8, y_tl8]).float().T)[: args.n_frames]

    # smooth the transition from intro to drop to prevent jerk
    translation.T[0, drop_start - args.fps : drop_start + args.fps] = ar.gaussian_filter(
        translation.T[0, drop_start - 5 * args.fps : drop_start + 5 * args.fps], 5
    )[4 * args.fps : -4 * args.fps]

    class Translate(NetworkBend):
        """From audioreactive/examples/bend.py"""

        def __init__(self, modulation, h, w, noise):
            sequential_fn = lambda b: th.nn.Sequential(
                th.nn.ReflectionPad2d((int(w / 2), int(w / 2), 0, 0)),  #  < Reflect out to 5x width (so that after
                th.nn.ReflectionPad2d((w, w, 0, 0)),  #                    < translating w pixels, center crop gives
                th.nn.ReflectionPad2d((w, 0, 0, 0)),  #                    < same features as translating 0 pixels)
                AddNoise(noise),  # add some noise to disguise reflections
                kT.Translate(b),
                kA.CenterCrop((h, w)),
            )
            super(Translate, self).__init__(sequential_fn, modulation)

    # create static noise for translate bend
    noise = 0.2 * th.randn((1, 1, h, 5 * w), device="cuda")
    # create function which returns an initialized Translate object when fed a batch of modulation
    # this is so that creation of the object is delayed until the specific batch is sent into the generator
    # (there's probably an easier way to do this without the kornia transforms, e.g. using Broad et al.'s transform implementations)
    transform = lambda batch: partial(Translate, h=h, w=w, noise=noise)(batch)
    bends += [{"layer": tl, "transform": transform, "modulation": translation}]  # add network bend to list dict

    return bends


================================================
FILE: audioreactive/examples/temper.py
================================================
"""
This file shows an example of spatial control of the noise using a simple circular mask
The latents are a chromagram weighted sequence, modulated by drum onsets
"""

import scipy.ndimage.filters as ndi
import torch as th

import audioreactive as ar

OVERRIDE = dict(audio_file="audioreactive/examples/Wavefunk - Temper.mp3", out_size=1024)


def initialize(args):
    # these onsets can definitely use some tweaking, the drum reactivity isn't great for this one
    # the main bass makes it hard to identify both the kick and the snare because it is so loud and covers the whole spectrum
    args.lo_onsets = ar.onsets(args.audio, args.sr, args.n_frames, fmax=150, smooth=5, clip=97, power=2)
    args.hi_onsets = ar.onsets(args.audio, args.sr, args.n_frames, fmin=500, smooth=5, clip=99, power=2)
    return args


def get_latents(selection, args):
    # create chromagram weighted sequence
    chroma = ar.chroma(args.audio, args.sr, args.n_frames)
    chroma_latents = ar.chroma_weight_latents(chroma, selection)
    latents = ar.gaussian_filter(chroma_latents, 4)

    # expand onsets to latent shape
    lo_onsets = args.lo_onsets[:, None, None]
    hi_onsets = args.hi_onsets[:, None, None]

    # modulate latents to specific latent vectors
    latents = hi_onsets * selection[[-4]] + (1 - hi_onsets) * latents
    latents = lo_onsets * selection[[-7]] + (1 - lo_onsets) * latents

    latents = ar.gaussian_filter(latents, 2, causal=0.2)

    return latents


def circular_mask(h, w, center=None, radius=None, soft=0):
    if center is None:  # use the middle of the image
        center = (int(w / 2), int(h / 2))
    if radius is None:  # use the smallest distance between the center and image walls
        radius = min(center[0], center[1], w - center[0], h - center[1])

    Y, X = np.ogrid[:h, :w]
    dist_from_center = np.sqrt((X - center[0]) ** 2 + (Y - center[1]) ** 2)
    mask = dist_from_center <= radius

    if soft > 0:
        mask = ndi.gaussian_filter(mask, sigma=int(round(soft)))  # blur mask for smoother transition

    return th.from_numpy(mask)


def get_noise(height, width, scale, num_scales, args):
    if width > 256:  # larger sizes don't fit in VRAM, just use default or randomize
        return None

    # expand onsets to noise shape
    # send to GPU as gaussian_filter on large noise tensors with high standard deviation is slow
    lo_onsets = args.lo_onsets[:, None, None, None].cuda()
    hi_onsets = args.hi_onsets[:, None, None, None].cuda()

    # 1s inside circle of radius, 0s outside
    mask = circular_mask(height, width, radius=int(width / 2), soft=2)[None, None, ...].float().cuda()

    # create noise which changes quickly (small standard deviation smoothing)
    noise_noisy = ar.gaussian_filter(th.randn((args.n_frames, 1, height, width), device="cuda"), 5)

    # create noise which changes slowly (large standard deviation smoothing)
    noise = ar.gaussian_filter(th.randn((args.n_frames, 1, height, width), device="cuda"), 128)

    # for lower layers, noise inside circle are affected by low onsets
    if width < 128:
        noise = 2 * mask * lo_onsets * noise_noisy + (1 - mask) * (1 - lo_onsets) * noise
    # for upper layers, noise outside circle are affected by high onsets
    if width > 32:
        noise = 0.75 * (1 - mask) * hi_onsets * noise_noisy + mask * (1 - 0.75 * hi_onsets) * noise

    # ensure amplitude of noise is close to standard normal distribution (dividing by std. dev. gets it exactly there)
    noise /= noise.std() * 2

    return noise.cpu()


================================================
FILE: audioreactive/latent.py
================================================
import gc

import numpy as np
import torch as th
from scipy import interpolate

from models.stylegan2 import Generator
from .signal import gaussian_filter

# ====================================================================================
# ================================= latent/noise ops =================================
# ====================================================================================


def chroma_weight_latents(chroma, latents):
    """Creates chromagram weighted latent sequence

    Args:
        chroma (th.tensor): Chromagram
        latents (th.tensor): Latents (must have same number as number of notes in chromagram)

    Returns:
        th.tensor: Chromagram weighted latent sequence
    """
    base_latents = (chroma[..., None, None] * latents[None, ...]).sum(1)
    return base_latents


def slerp(val, low, high):
    """Interpolation along geodesic of n-dimensional unit sphere
    from https://github.com/soumith/dcgan.torch/issues/14#issuecomment-200025792

    Args:
        val (float): Value between 0 and 1 representing fraction of interpolation completed
        low (float): Starting value
        high (float): Ending value

    Returns:
        float: Interpolated value
    """
    omega = np.arccos(np.clip(np.dot(low / np.linalg.norm(low), high / np.linalg.norm(high)), -1, 1))
    so = np.sin(omega)
    if so == 0:
        return (1.0 - val) * low + val * high  # L'Hopital's rule/LERP
    return np.sin((1.0 - val) * omega) / so * low + np.sin(val * omega) / so * high


def slerp_loops(latent_selection, n_frames, n_loops, smoothing=1, loop=True):
    """Get looping latents using geodesic interpolation. Total length of n_frames with n_loops repeats.

    Args:
        latent_selection (th.tensor): Set of latents to loop between (in order)
        n_frames (int): Total length of output looping sequence
        n_loops (int): Number of times to loop
        smoothing (int, optional): Standard deviation of gaussian smoothing kernel. Defaults to 1.
        loop (bool, optional): Whether to return to first latent. Defaults to True.

    Returns:
        th.tensor: Sequence of smoothly looping latents
    """
    if loop:
        latent_selection = np.concatenate([latent_selection, latent_selection[[0]]])

    base_latents = []
    for n in range(len(latent_selection)):
        for val in np.linspace(0.0, 1.0, int(n_frames // max(1, n_loops) // len(latent_selection))):
            base_latents.append(
                th.from_numpy(
                    slerp(
                        val,
                        latent_selection[n % len(latent_selection)][0],
                        latent_selection[(n + 1) % len(latent_selection)][0],
                    )
                )
            )
    base_latents = th.stack(base_latents)
    base_latents = gaussian_filter(base_latents, smoothing)
    base_latents = th.cat([base_latents] * int(n_frames / len(base_latents)), axis=0)
    base_latents = th.cat([base_latents[:, None, :]] * 18, axis=1)
    if n_frames - len(base_latents) != 0:
        base_latents = th.cat([base_latents, base_latents[0 : n_frames - len(base_latents)]])
    return base_latents


def spline_loops(latent_selection, n_frames, n_loops, loop=True):
    """Get looping latents using spline interpolation. Total length of n_frames with n_loops repeats.

    Args:
        latent_selection (th.tensor): Set of latents to loop between (in order)
        n_frames (int): Total length of output looping sequence
        n_loops (int): Number of times to loop
        loop (bool, optional): Whether to return to first latent. Defaults to True.

    Returns:
        th.tensor: Sequence of smoothly looping latents
    """
    if loop:
        latent_selection = np.concatenate([latent_selection, latent_selection[[0]]])

    x = np.linspace(0, 1, int(n_frames // max(1, n_loops)))
    base_latents = np.zeros((len(x), *latent_selection.shape[1:]))
    for lay in range(latent_selection.shape[1]):
        for lat in range(latent_selection.shape[2]):
            tck = interpolate.splrep(np.linspace(0, 1, latent_selection.shape[0]), latent_selection[:, lay, lat])
            base_latents[:, lay, lat] = interpolate.splev(x, tck)

    base_latents = th.cat([th.from_numpy(base_latents)] * int(n_frames / len(base_latents)), axis=0)
    if n_frames - len(base_latents) > 0:
        base_latents = th.cat([base_latents, base_latents[0 : n_frames - len(base_latents)]])
    return base_latents[:n_frames]


def wrapping_slice(tensor, start, length, return_indices=False):
    """Gets slice of tensor of a given length that wraps around to beginning

    Args:
        tensor (th.tensor): Tensor to slice
        start (int): Starting index
        length (int): Size of slice
        return_indices (bool, optional): Whether to return indices rather than values. Defaults to False.

    Returns:
        th.tensor: Values or indices of slice
    """
    if start + length <= tensor.shape[0]:
        indices = th.arange(start, start + length)
    else:
        indices = th.cat((th.arange(start, tensor.shape[0]), th.arange(0, (start + length) % tensor.shape[0])))
    if tensor.shape[0] == 1:
        indices = th.zeros(1, dtype=th.int64)
    if return_indices:
        return indices
    return tensor[indices]


def generate_latents(n_latents, ckpt, G_res, noconst=False, latent_dim=512, n_mlp=8, channel_multiplier=2):
    """Generates random, mapped latents

    Args:
        n_latents (int): Number of mapped latents to generate 
        ckpt (str): Generator checkpoint to use
        G_res (int): Generator's training resolution
        noconst (bool, optional): Whether the generator was trained without constant starting layer. Defaults to False.
        latent_dim (int, optional): Size of generator's latent vectors. Defaults to 512.
        n_mlp (int, optional): Number of layers in the generator's mapping network. Defaults to 8.
        channel_multiplier (int, optional): Scaling multiplier for generator's channel depth. Defaults to 2.

    Returns:
        th.tensor: Set of mapped latents
    """
    generator = Generator(
        G_res, latent_dim, n_mlp, channel_multiplier=channel_multiplier, constant_input=not noconst, checkpoint=ckpt,
    ).cuda()
    zs = th.randn((n_latents, latent_dim), device="cuda")
    latent_selection = generator(zs, map_latents=True).cpu()
    del generator, zs
    gc.collect()
    th.cuda.empty_cache()
    return latent_selection


def save_latents(latents, filename):
    """Saves latent vectors to file

    Args:
        latents (th.tensor): Latent vector(s) to save
        filename (str): Filename to save to
    """
    np.save(filename, latents)


def load_latents(filename):
    """Load latents from numpy file

    Args:
        filename (str): Filename to load from

    Returns:
        th.tensor: Latent vectors
    """
    return th.from_numpy(np.load(filename))


def _perlinterpolant(t):
    return t * t * t * (t * (t * 6 - 15) + 10)


def perlin_noise(shape, res, tileable=(True, False, False), interpolant=_perlinterpolant):
    """Generate a 3D tensor of perlin noise.

    Args:
        shape: The shape of the generated tensor (tuple of three ints). This must be a multiple of res.
        res: The number of periods of noise to generate along each axis (tuple of three ints). Note shape must be a multiple of res.
        tileable: If the noise should be tileable along each axis (tuple of three bools). Defaults to (False, False, False).
        interpolant: The interpolation function, defaults to t*t*t*(t*(t*6 - 15) + 10).

    Returns:
        A tensor of shape shape with the generated noise.

    Raises:
        ValueError: If shape is not a multiple of res.
    """
    delta = (res[0] / shape[0], res[1] / shape[1], res[2] / shape[2])
    d = (shape[0] // res[0], shape[1] // res[1], shape[2] // res[2])
    grid = np.mgrid[0 : res[0] : delta[0], 0 : res[1] : delta[1], 0 : res[2] : delta[2]]
    grid = grid.transpose(1, 2, 3, 0) % 1
    grid = th.from_numpy(grid).cuda()
    # Gradients
    theta = 2 * np.pi * np.random.rand(res[0] + 1, res[1] + 1, res[2] + 1)
    phi = 2 * np.pi * np.random.rand(res[0] + 1, res[1] + 1, res[2] + 1)
    gradients = np.stack((np.sin(phi) * np.cos(theta), np.sin(phi) * np.sin(theta), np.cos(phi)), axis=3)
    if tileable[0]:
        gradients[-1, :, :] = gradients[0, :, :]
    if tileable[1]:
        gradients[:, -1, :] = gradients[:, 0, :]
    if tileable[2]:
        gradients[:, :, -1] = gradients[:, :, 0]
    gradients = gradients.repeat(d[0], 0).repeat(d[1], 1).repeat(d[2], 2)
    gradients = th.from_numpy(gradients).cuda()
    g000 = gradients[: -d[0], : -d[1], : -d[2]]
    g100 = gradients[d[0] :, : -d[1], : -d[2]]
    g010 = gradients[: -d[0], d[1] :, : -d[2]]
    g110 = gradients[d[0] :, d[1] :, : -d[2]]
    g001 = gradients[: -d[0], : -d[1], d[2] :]
    g101 = gradients[d[0] :, : -d[1], d[2] :]
    g011 = gradients[: -d[0], d[1] :, d[2] :]
    g111 = gradients[d[0] :, d[1] :, d[2] :]
    # Ramps
    n000 = th.sum(th.stack((grid[:, :, :, 0], grid[:, :, :, 1], grid[:, :, :, 2]), axis=3) * g000, 3)
    n100 = th.sum(th.stack((grid[:, :, :, 0] - 1, grid[:, :, :, 1], grid[:, :, :, 2]), axis=3) * g100, 3)
    n010 = th.sum(th.stack((grid[:, :, :, 0], grid[:, :, :, 1] - 1, grid[:, :, :, 2]), axis=3) * g010, 3)
    n110 = th.sum(th.stack((grid[:, :, :, 0] - 1, grid[:, :, :, 1] - 1, grid[:, :, :, 2]), axis=3) * g110, 3)
    n001 = th.sum(th.stack((grid[:, :, :, 0], grid[:, :, :, 1], grid[:, :, :, 2] - 1), axis=3) * g001, 3)
    n101 = th.sum(th.stack((grid[:, :, :, 0] - 1, grid[:, :, :, 1], grid[:, :, :, 2] - 1), axis=3) * g101, 3)
    n011 = th.sum(th.stack((grid[:, :, :, 0], grid[:, :, :, 1] - 1, grid[:, :, :, 2] - 1), axis=3) * g011, 3)
    n111 = th.sum(th.stack((grid[:, :, :, 0] - 1, grid[:, :, :, 1] - 1, grid[:, :, :, 2] - 1), axis=3) * g111, 3)
    # Interpolation
    t = interpolant(grid)
    n00 = n000 * (1 - t[:, :, :, 0]) + t[:, :, :, 0] * n100
    n10 = n010 * (1 - t[:, :, :, 0]) + t[:, :, :, 0] * n110
    n01 = n001 * (1 - t[:, :, :, 0]) + t[:, :, :, 0] * n101
    n11 = n011 * (1 - t[:, :, :, 0]) + t[:, :, :, 0] * n111
    n0 = (1 - t[:, :, :, 1]) * n00 + t[:, :, :, 1] * n10
    n1 = (1 - t[:, :, :, 1]) * n01 + t[:, :, :, 1] * n11
    perlin = (1 - t[:, :, :, 2]) * n0 + t[:, :, :, 2] * n1
    return perlin * 2 - 1  # stretch from -1 to 1


================================================
FILE: audioreactive/signal.py
================================================
import os
import warnings
from pathlib import Path

import joblib
import librosa as rosa
import librosa.display
import madmom as mm
import matplotlib.patches as patches
import matplotlib.pyplot as plt
import numpy as np
import scipy
import scipy.signal as signal
import sklearn.cluster
import torch as th
import torch.nn.functional as F

SMF = 1  # this is set by generate_audiovisual.py based on rendering fps


def set_SMF(smf):
    global SMF
    SMF = smf


# ====================================================================================
# ==================================== signal ops ====================================
# ====================================================================================


def onsets(audio, sr, n_frames, margin=8, fmin=20, fmax=8000, smooth=1, clip=100, power=1, type="mm"):
    """Creates onset envelope from audio

    Args:
        audio (np.array): Audio signal
        sr (int): Sampling rate of the audio
        n_frames (int): Total number of frames to resample envelope to
        margin (int, optional): For percussive source separation, higher values create more extreme separations. Defaults to 8.
        fmin (int, optional): Minimum frequency for onset analysis. Defaults to 20.
        fmax (int, optional): Maximum frequency for onset analysis. Defaults to 8000.
        smooth (int, optional): Standard deviation of gaussian kernel to smooth with. Defaults to 1.
        clip (int, optional): Percentile to clip onset signal to. Defaults to 100.
        power (int, optional): Exponent to raise onset signal to. Defaults to 1.
        type (str, optional): ["rosa", "mm"]. Whether to use librosa or madmom for onset analysis. Madmom is slower but often more accurate. Defaults to "mm".

    Returns:
        th.tensor, shape=(n_frames,): Onset envelope
    """
    y_perc = rosa.effects.percussive(y=audio, margin=margin)
    if type == "rosa":
        onset = rosa.onset.onset_strength(y=y_perc, sr=sr, fmin=fmin, fmax=fmax)
    elif type == "mm":
        sig = mm.audio.signal.Signal(y_perc, num_channels=1, sample_rate=sr)
        sig_frames = mm.audio.signal.FramedSignal(sig, frame_size=2048, hop_size=441)
        stft = mm.audio.stft.ShortTimeFourierTransform(sig_frames, circular_shift=True)
        spec = mm.audio.spectrogram.Spectrogram(stft, circular_shift=True)
        filt_spec = mm.audio.spectrogram.FilteredSpectrogram(spec, num_bands=24, fmin=fmin, fmax=fmax)
        onset = np.sum(
            [
                mm.features.onsets.spectral_diff(filt_spec),
                mm.features.onsets.spectral_flux(filt_spec),
                mm.features.onsets.superflux(filt_spec),
                mm.features.onsets.complex_flux(filt_spec),
                mm.features.onsets.modified_kullback_leibler(filt_spec),
            ],
            axis=0,
        )
    onset = np.clip(signal.resample(onset, n_frames), onset.min(), onset.max())
    onset = th.from_numpy(onset).float()
    onset = gaussian_filter(onset, smooth, causal=0)
    onset = percentile_clip(onset, clip)
    onset = onset ** power
    return onset


def rms(y, sr, n_frames, fmin=20, fmax=8000, smooth=180, clip=50, power=6):
    """Creates RMS envelope from audio

    Args:
        audio (np.array): Audio signal
        sr (int): Sampling rate of the audio
        n_frames (int): Total number of frames to resample envelope to
        fmin (int, optional): Minimum frequency for onset analysis. Defaults to 20.
        fmax (int, optional): Maximum frequency for onset analysis. Defaults to 8000.
        smooth (int, optional): Standard deviation of gaussian kernel to smooth with. Defaults to 180.
        clip (int, optional): Percentile to clip onset signal to. Defaults to 50.
        power (int, optional): Exponent to raise onset signal to. Defaults to 6.

    Returns:
        th.tensor, shape=(n_frames,): RMS envelope
    """
    y_filt = signal.sosfilt(signal.butter(12, [fmin, fmax], "bp", fs=sr, output="sos"), y)
    rms = rosa.feature.rms(S=np.abs(rosa.stft(y=y_filt, hop_length=512)))[0]
    rms = np.clip(signal.resample(rms, n_frames), rms.min(), rms.max())
    rms = th.from_numpy(rms).float()
    rms = gaussian_filter(rms, smooth, causal=0.05)
    rms = percentile_clip(rms, clip)
    rms = rms ** power
    return rms


def raw_chroma(audio, sr, type="cens", nearest_neighbor=True):
    """Creates chromagram

    Args:
        audio (np.array): Audio signal
        sr (int): Sampling rate of the audio
        type (str, optional): ["cens", "cqt", "stft", "deep", "clp"]. Which strategy to use to calculate the chromagram. Defaults to "cens".
        nearest_neighbor (bool, optional): Whether to post process using nearest neighbor smoothing. Defaults to True.

    Returns:
        np.array, shape=(12, n_frames): Chromagram
    """
    if type == "cens":
        ch = rosa.feature.chroma_cens(y=audio, sr=sr)
    elif type == "cqt":
        ch = rosa.feature.chroma_cqt(y=audio, sr=sr)
    elif type == "stft":
        ch = rosa.feature.chroma_stft(y=audio, sr=sr)
    elif type == "deep":
        sig = mm.audio.signal.Signal(audio, num_channels=1, sample_rate=sr)
        ch = mm.audio.chroma.DeepChromaProcessor().process(sig).T
    elif type == "clp":
        sig = mm.audio.signal.Signal(audio, num_channels=1, sample_rate=sr)
        ch = mm.audio.chroma.CLPChromaProcessor().process(sig).T
    else:
        print("chroma type not recognized, options are: [cens, cqt, deep, clp, or stft]. defaulting to cens...")
        ch = rosa.feature.chroma_cens(y=audio, sr=sr)

    if nearest_neighbor:
        ch = np.minimum(ch, rosa.decompose.nn_filter(ch, aggregate=np.median, metric="cosine"))

    return ch


def chroma(audio, sr, n_frames, margin=16, type="cens", notes=12):
    """Creates chromagram for the harmonic component of the audio

    Args:
        audio (np.array): Audio signal
        sr (int): Sampling rate of the audio
        n_frames (int): Total number of frames to resample envelope to
        margin (int, optional): For harmonic source separation, higher values create more extreme separations. Defaults to 16.
        type (str, optional): ["cens", "cqt", "stft", "deep", "clp"]. Which strategy to use to calculate the chromagram. Defaults to "cens".
        notes (int, optional): Number of notes to use in output chromagram (e.g. 5 for pentatonic scale, 7 for standard western scales). Defaults to 12.

    Returns:
        th.tensor, shape=(n_frames, 12): Chromagram
    """
    y_harm = rosa.effects.harmonic(y=audio, margin=margin)
    chroma = raw_chroma(y_harm, sr, type=type).T
    chroma = signal.resample(chroma, n_frames)
    notes_indices = np.argsort(np.median(chroma, axis=0))[:notes]
    chroma = chroma[:, notes_indices]
    chroma = th.from_numpy(chroma / chroma.sum(1)[:, None]).float()
    return chroma


def laplacian_segmentation(signal, sr, k=5, plot=False):
    """Segments the audio with pattern recurrence analysis
    From https://librosa.org/doc/latest/auto_examples/plot_segmentation.html#sphx-glr-auto-examples-plot-segmentation-py%22

    Args:
        signal (np.array): Audio signal
        sr (int): Sampling rate of the audio
        k (int, optional): Number of labels to use during segmentation. Defaults to 5.
        plot (bool, optional): Whether to show plot of found segmentation. Defaults to False.

    Returns:
        tuple(list, list): List of starting timestamps and labels of found segments
    """
    BINS_PER_OCTAVE = 12 * 3
    N_OCTAVES = 7
    C = librosa.amplitude_to_db(
        np.abs(librosa.cqt(y=signal, sr=sr, bins_per_octave=BINS_PER_OCTAVE, n_bins=N_OCTAVES * BINS_PER_OCTAVE)),
        ref=np.max,
    )

    # make CQT beat-synchronous to reduce dimensionality
    tempo, beats = librosa.beat.beat_track(y=signal, sr=sr, trim=False)
    Csync = librosa.util.sync(C, beats, aggregate=np.median)

    # build a weighted recurrence matrix using beat-synchronous CQT
    R = librosa.segment.recurrence_matrix(Csync, width=3, mode="affinity", sym=True)
    # enhance diagonals with a median filter
    df = librosa.segment.timelag_filter(scipy.ndimage.median_filter)
    Rf = df(R, size=(1, 7))

    # build the sequence matrix using mfcc-similarity
    mfcc = librosa.feature.mfcc(y=signal, sr=sr)
    Msync = librosa.util.sync(mfcc, beats)
    path_distance = np.sum(np.diff(Msync, axis=1) ** 2, axis=0)
    sigma = np.median(path_distance)
    path_sim = np.exp(-path_distance / sigma)
    R_path = np.diag(path_sim, k=1) + np.diag(path_sim, k=-1)

    # compute the balanced combination
    deg_path = np.sum(R_path, axis=1)
    deg_rec = np.sum(Rf, axis=1)
    mu = deg_path.dot(deg_path + deg_rec) / np.sum((deg_path + deg_rec) ** 2)

    A = mu * Rf + (1 - mu) * R_path
    # compute the normalized laplacian and its spectral decomposition
    L = scipy.sparse.csgraph.laplacian(A, normed=True)
    evals, evecs = scipy.linalg.eigh(L)
    # median filter to smooth over small discontinuities
    evecs = scipy.ndimage.median_filter(evecs, size=(9, 1))
    # cumulative normalization for symmetric normalized laplacian eigenvectors
    Cnorm = np.cumsum(evecs ** 2, axis=1) ** 0.5

    X = evecs[:, :k] / Cnorm[:, k - 1 : k]

    # use first k components to cluster beats into segments
    seg_ids = sklearn.cluster.KMeans(n_clusters=k).fit_predict(X)

    bound_beats = 1 + np.flatnonzero(seg_ids[:-1] != seg_ids[1:])  # locate segment boundaries from the label sequence
    bound_beats = librosa.util.fix_frames(bound_beats, x_min=0)  # count beat 0 as a boundary
    bound_segs = list(seg_ids[bound_beats])  # compute the segment label for each boundary
    bound_frames = beats[bound_beats]  # convert beat indices to frames
    bound_frames = librosa.util.fix_frames(bound_frames, x_min=None, x_max=C.shape[1] - 1)
    bound_times = librosa.frames_to_time(bound_frames)
    if bound_times[0] != 0:
        bound_times[0] = 0

    if plot:
        freqs = librosa.cqt_frequencies(
            n_bins=C.shape[0], fmin=librosa.note_to_hz("C1"), bins_per_octave=BINS_PER_OCTAVE
        )
        fig, ax = plt.subplots()
        colors = plt.get_cmap("Paired", k)
        librosa.display.specshow(C, y_axis="cqt_hz", sr=sr, bins_per_octave=BINS_PER_OCTAVE, x_axis="time", ax=ax)
        for interval, label in zip(zip(bound_times, bound_times[1:]), bound_segs):
            ax.add_patch(
                patches.Rectangle(
                    (interval[0], freqs[0]), interval[1] - interval[0], freqs[-1], facecolor=colors(label), alpha=0.50
                )
            )
        plt.show()

    return list(bound_times), list(bound_segs)


def normalize(signal):
    """Normalize signal between 0 and 1

    Args:
        signal (np.array/th.tensor): Signal to normalize

    Returns:
        np.array/th.tensor: Normalized signal
    """
    signal -= signal.min()
    signal /= signal.max()
    return signal


def percentile(signal, p):
    """Calculate percentile of signal

    Args:
        signal (np.array/th.tensor): Signal to normalize
        p (int): [0-100]. Percentile to find

    Returns:
        int: Percentile signal value
    """
    k = 1 + round(0.01 * float(p) * (signal.numel() - 1))
    return signal.view(-1).kthvalue(k).values.item()


def percentile_clip(signal, p):
    """Normalize signal between 0 and 1, clipping peak values above given percentile

    Args:
        signal (th.tensor): Signal to normalize
        p (int): [0-100]. Percentile to clip to

    Returns:
        th.tensor: Clipped signal
    """
    locs = th.arange(0, signal.shape[0])
    peaks = th.ones(signal.shape, dtype=bool)
    main = signal.take(locs)

    plus = signal.take((locs + 1).clamp(0, signal.shape[0] - 1))
    minus = signal.take((locs - 1).clamp(0, signal.shape[0] - 1))
    peaks &= th.gt(main, plus)
    peaks &= th.gt(main, minus)

    signal = signal.clamp(0, percentile(signal[peaks], p))
    signal /= signal.max()
    return signal


def compress(signal, threshold, ratio, invert=False):
    """Expand or compress signal. Values above/below (depending on invert) threshold are multiplied by ratio.

    Args:
        signal (th.tensor): Signal to normalize
        threshold (int): Signal value above/below which to change signal
        ratio (float): Value to multiply signal with
        invert (bool, optional): Specifies if values above or below threshold are affected. Defaults to False.

    Returns:
        th.tensor: Compressed/expanded signal
    """
    if invert:
        signal[signal < threshold] *= ratio
    else:
        signal[signal > threshold] *= ratio
    return normalize(signal)


def expand(signal, threshold, ratio, invert=False):
    """Alias for compress. Whether compression or expansion occurs is determined by values of threshold and ratio"""
    return compress(signal, threshold, ratio, invert)


def gaussian_filter(x, sigma, causal=None):
    """Smooth tensors along time (first) axis with gaussian kernel.

    Args:
        x (th.tensor): Tensor to be smoothed
        sigma (float): Standard deviation for gaussian kernel (higher value gives smoother result)
        causal (float, optional): Factor to multiply right side of gaussian kernel with. Lower value decreases effect of "future" values. Defaults to None.

    Returns:
        th.tensor: Smoothed tensor
    """
    dim = len(x.shape)
    n_frames = x.shape[0]
    while len(x.shape) < 3:
        x = x[:, None]

    radius = min(int(sigma * 4 * SMF), 3 * len(x))
    channels = x.shape[1]

    kernel = th.arange(-radius, radius + 1, dtype=th.float32, device=x.device)
    kernel = th.exp(-0.5 / sigma ** 2 * kernel ** 2)
    if causal is not None:
        kernel[radius + 1 :] *= 0 if not isinstance(causal, float) else causal
    kernel = kernel / kernel.sum()
    kernel = kernel.view(1, 1, len(kernel)).repeat(channels, 1, 1)

    if dim == 4:
        t, c, h, w = x.shape
        x = x.view(t, c, h * w)
    x = x.transpose(0, 2)

    if radius > n_frames:  # prevent padding errors on short sequences
        x = F.pad(x, (n_frames, n_frames), mode="circular")
        print(
            f"WARNING: Gaussian filter radius ({int(sigma * 4 * SMF)}) is larger than number of frames ({n_frames}).\n\t Filter size has been lowered to ({radius}). You might want to consider lowering sigma ({sigma})."
        )
        x = F.pad(x, (radius - n_frames, radius - n_frames), mode="constant")
    else:
        x = F.pad(x, (radius, radius), mode="circular")

    x = F.conv1d(x, weight=kernel, groups=channels)

    x = x.transpose(0, 2)
    if dim == 4:
        x = x.view(t, c, h, w)

    if len(x.shape) > dim:
        x = x.squeeze()

    return x


def load_audio(audio_file, offset=0, duration=-1, cache=True):
    """Handles loading of audio files. Automatically caches to .npy files to increase loading speed.

    Args:
        audio_file (str): Path to audio file to load
        offset (float, optional): Time (in seconds) to start from. Defaults to 0.
        duration (float, optional): Length of time to load in. Defaults to -1 (full duration).
        cache (bool): Whether to cache the raw audio file or not

    Returns:
        audio   : audio signal
        sr      : sample rate of audio
        duration: duration of audio in seconds
    """
    audio_dur = rosa.get_duration(filename=audio_file)
    if duration == -1 or audio_dur < duration:
        duration = audio_dur
        if offset != 0:
            duration -= offset

    cache_file = (
        f"workspace/{Path(audio_file).stem}"
        + ("" if duration == -1 else f"_length{duration}")
        + ("" if offset == 0 else f"_start{offset}")
        + ".npy"
    )
    if cache and not os.path.exists(cache_file):
        with warnings.catch_warnings():
            warnings.filterwarnings("ignore", message="PySoundFile failed. Trying audioread instead.")
            audio, sr = rosa.load(audio_file, offset=offset, duration=duration)
        joblib.dump((audio, sr), cache_file)
    else:
        audio, sr = joblib.load(cache_file)

    return audio, sr, duration


================================================
FILE: audioreactive/util.py
================================================
import librosa as rosa
import librosa.display
import matplotlib.pyplot as plt
import numpy as np

# ====================================================================================
# ==================================== utilities =====================================
# ====================================================================================


def info(arr):
    """Shows statistics and shape information of (lists of) np.arrays/th.tensors

    Args:
        arr (np.array/th.tensor/list): List of or single np.array or th.tensor
    """
    if isinstance(arr, list):
        print([(list(a.shape), f"{a.min():.2f}", f"{a.mean():.2f}", f"{a.max():.2f}") for a in arr])
    else:
        print(list(arr.shape), f"{arr.min():.2f}", f"{arr.mean():.2f}", f"{arr.max():.2f}")


def plot_signals(signals):
    """Shows plot of (multiple) 1D signals

    Args:
        signals (np.array/th.tensor): List of signals (1 non-unit dimension)
    """
    plt.figure(figsize=(16, 4 * len(signals)))
    for sbplt, y in enumerate(signals):
        try:
            signal = signal.cpu().numpy()
        except:
            pass
        plt.subplot(len(signals), 1, sbplt + 1)
        plt.plot(y.squeeze())
    plt.tight_layout()
    plt.show()


def plot_spectra(spectra, chroma=False):
    """Shows plot of (multiple) spectrograms

    Args:
        spectra (np.array/th.tensor): List of spectrograms
        chroma (bool, optional): Whether to plot with chromagram y-axis label. Defaults to False.
    """
    fig, axes = plt.subplots(len(spectra), 1, figsize=(16, 4 * len(spectra)))
    for ax, spectrum in zip(axes if len(spectra) > 1 else [axes], spectra):
        try:
            spectrum = spectrum.cpu().numpy()
        except:
            pass
        if spectrum.shape[1] == 12:
            spectrum = spectrum.T
        rosa.display.specshow(spectrum, y_axis="chroma" if chroma else None, x_axis="time", ax=ax)
    plt.tight_layout()
    plt.show()


def plot_audio(audio, sr):
    """Shows spectrogram of audio signal

    Args:
        audio (np.array): Audio signal to be plotted
        sr (int): Sampling rate of the audio
    """
    plt.figure(figsize=(16, 9))
    rosa.display.specshow(
        rosa.power_to_db(rosa.feature.melspectrogram(y=audio, sr=sr), ref=np.max), y_axis="mel", x_axis="time"
    )
    plt.colorbar(format="%+2.f dB")
    plt.tight_layout()
    plt.show()


def plot_chroma_comparison(audio, sr):
    """Shows plot comparing different chromagram strategies.

    Args:
        audio (np.array): Audio signal to be plotted
        sr (int): Sampling rate of the audio
    """
    fig, ax = plt.subplots(nrows=2, ncols=3, figsize=(16, 9))
    for col, types in enumerate([["cens", "cqt"], ["deep", "clp"], ["stft"]]):
        for row, type in enumerate(types):
            ch = raw_chroma(audio, sr, type=type)
            if ch.shape[1] == 12:
                ch = ch.T
            librosa.display.specshow(ch, y_axis="chroma", x_axis="time", ax=ax[row, col])
            ax[row, col].set(title=type)
            ax[row, col].label_outer()
    plt.tight_layout()
    plt.show()


================================================
FILE: augment.py
================================================
import math

import torch
from torch.nn import functional as F

from op import upfirdn2d


SYM6 = (
    0.015404109327027373,
    0.0034907120842174702,
    -0.11799011114819057,
    -0.048311742585633,
    0.4910559419267466,
    0.787641141030194,
    0.3379294217276218,
    -0.07263752278646252,
    -0.021060292512300564,
    0.04472490177066578,
    0.0017677118642428036,
    -0.007800708325034148,
)


def translate_mat(t_x, t_y):
    batch = t_x.shape[0]

    mat = torch.eye(3).unsqueeze(0).repeat(batch, 1, 1)
    translate = torch.stack((t_x, t_y), 1)
    mat[:, :2, 2] = translate

    return mat


def rotate_mat(theta):
    batch = theta.shape[0]

    mat = torch.eye(3).unsqueeze(0).repeat(batch, 1, 1)
    sin_t = torch.sin(theta)
    cos_t = torch.cos(theta)
    rot = torch.stack((cos_t, -sin_t, sin_t, cos_t), 1).view(batch, 2, 2)
    mat[:, :2, :2] = rot

    return mat


def scale_mat(s_x, s_y):
    batch = s_x.shape[0]

    mat = torch.eye(3).unsqueeze(0).repeat(batch, 1, 1)
    mat[:, 0, 0] = s_x
    mat[:, 1, 1] = s_y

    return mat


def translate3d_mat(t_x, t_y, t_z):
    batch = t_x.shape[0]

    mat = torch.eye(4).unsqueeze(0).repeat(batch, 1, 1)
    translate = torch.stack((t_x, t_y, t_z), 1)
    mat[:, :3, 3] = translate

    return mat


def rotate3d_mat(axis, theta):
    batch = theta.shape[0]

    u_x, u_y, u_z = axis

    eye = torch.eye(3).unsqueeze(0)
    cross = torch.tensor([(0, -u_z, u_y), (u_z, 0, -u_x), (-u_y, u_x, 0)]).unsqueeze(0)
    outer = torch.tensor(axis)
    outer = (outer.unsqueeze(1) * outer).unsqueeze(0)

    sin_t = torch.sin(theta).view(-1, 1, 1)
    cos_t = torch.cos(theta).view(-1, 1, 1)

    rot = cos_t * eye + sin_t * cross + (1 - cos_t) * outer

    eye_4 = torch.eye(4).unsqueeze(0).repeat(batch, 1, 1)
    eye_4[:, :3, :3] = rot

    return eye_4


def scale3d_mat(s_x, s_y, s_z):
    batch = s_x.shape[0]

    mat = torch.eye(4).unsqueeze(0).repeat(batch, 1, 1)
    mat[:, 0, 0] = s_x
    mat[:, 1, 1] = s_y
    mat[:, 2, 2] = s_z

    return mat


def luma_flip_mat(axis, i):
    batch = i.shape[0]

    eye = torch.eye(4).unsqueeze(0).repeat(batch, 1, 1)
    axis = torch.tensor(axis + (0,))
    flip = 2 * torch.ger(axis, axis) * i.view(-1, 1, 1)

    return eye - flip


def saturation_mat(axis, i):
    batch = i.shape[0]

    eye = torch.eye(4).unsqueeze(0).repeat(batch, 1, 1)
    axis = torch.tensor(axis + (0,))
    axis = torch.ger(axis, axis)
    saturate = axis + (eye - axis) * i.view(-1, 1, 1)

    return saturate


def lognormal_sample(size, mean=0, std=1):
    return torch.empty(size).log_normal_(mean=mean, std=std)


def category_sample(size, categories):
    category = torch.tensor(categories)
    sample = torch.randint(high=len(categories), size=(size,))

    return category[sample]


def uniform_sample(size, low, high):
    return torch.empty(size).uniform_(low, high)


def normal_sample(size, mean=0, std=1):
    return torch.empty(size).normal_(mean, std)


def bernoulli_sample(size, p):
    return torch.empty(size).bernoulli_(p)


def random_mat_apply(p, transform, prev, eye):
    size = transform.shape[0]
    select = bernoulli_sample(size, p).view(size, 1, 1)
    select_transform = select * transform + (1 - select) * eye

    return select_transform @ prev


def sample_affine(p, size, height, width):
    G = torch.eye(3).unsqueeze(0).repeat(size, 1, 1)
    eye = G

    # flip
    param = category_sample(size, (0, 1))
    Gc = scale_mat(1 - 2.0 * param, torch.ones(size))
    G = random_mat_apply(p, Gc, G, eye)
    # print('flip', G, scale_mat(1 - 2.0 * param, torch.ones(size)), sep='\n')

    # 90 rotate
    param = category_sample(size, (0, 3))
    Gc = rotate_mat(-math.pi / 2 * param)
    G = random_mat_apply(p, Gc, G, eye)
    # print('90 rotate', G, rotate_mat(-math.pi / 2 * param), sep='\n')

    # integer translate
    param = uniform_sample(size, -0.125, 0.125)
    param_height = torch.round(param * height) / height
    param_width = torch.round(param * width) / width
    Gc = translate_mat(param_width, param_height)
    G = random_mat_apply(p, Gc, G, eye)
    # print('integer translate', G, translate_mat(param_width, param_height), sep='\n')

    # isotropic scale
    param = lognormal_sample(size, std=0.2 * math.log(2))
    Gc = scale_mat(param, param)
    G = random_mat_apply(p, Gc, G, eye)
    # print('isotropic scale', G, scale_mat(param, param), sep='\n')

    p_rot = 1 - math.sqrt(1 - p)

    # pre-rotate
    param = uniform_sample(size, -math.pi, math.pi)
    Gc = rotate_mat(-param)
    G = random_mat_apply(p_rot, Gc, G, eye)
    # print('pre-rotate', G, rotate_mat(-param), sep='\n')

    # anisotropic scale
    param = lognormal_sample(size, std=0.2 * math.log(2))
    Gc = scale_mat(param, 1 / param)
    G = random_mat_apply(p, Gc, G, eye)
    # print('anisotropic scale', G, scale_mat(param, 1 / param), sep='\n')

    # post-rotate
    param = uniform_sample(size, -math.pi, math.pi)
    Gc = rotate_mat(-param)
    G = random_mat_apply(p_rot, Gc, G, eye)
    # print('post-rotate', G, rotate_mat(-param), sep='\n')

    # fractional translate
    param = normal_sample(size, std=0.125)
    Gc = translate_mat(param, param)
    G = random_mat_apply(p, Gc, G, eye)
    # print('fractional translate', G, translate_mat(param, param), sep='\n')

    return G


def sample_color(p, size):
    C = torch.eye(4).unsqueeze(0).repeat(size, 1, 1)
    eye = C
    axis_val = 1 / math.sqrt(3)
    axis = (axis_val, axis_val, axis_val)

    # brightness
    param = normal_sample(size, std=0.2)
    Cc = translate3d_mat(param, param, param)
    C = random_mat_apply(p, Cc, C, eye)

    # contrast
    param = lognormal_sample(size, std=0.5 * math.log(2))
    Cc = scale3d_mat(param, param, param)
    C = random_mat_apply(p, Cc, C, eye)

    # luma flip
    param = category_sample(size, (0, 1))
    Cc = luma_flip_mat(axis, param)
    C = random_mat_apply(p, Cc, C, eye)

    # hue rotation
    param = uniform_sample(size, -math.pi, math.pi)
    Cc = rotate3d_mat(axis, param)
    C = random_mat_apply(p, Cc, C, eye)

    # saturation
    param = lognormal_sample(size, std=1 * math.log(2))
    Cc = saturation_mat(axis, param)
    C = random_mat_apply(p, Cc, C, eye)

    return C


def make_grid(shape, x0, x1, y0, y1, device):
    n, c, h, w = shape
    grid = torch.empty(n, h, w, 3, device=device)
    grid[:, :, :, 0] = torch.linspace(x0, x1, w, device=device)
    grid[:, :, :, 1] = torch.linspace(y0, y1, h, device=device).unsqueeze(-1)
    grid[:, :, :, 2] = 1

    return grid


def affine_grid(grid, mat):
    n, h, w, _ = grid.shape
    return (grid.view(n, h * w, 3) @ mat.transpose(1, 2)).view(n, h, w, 2)


def get_padding(G, height, width):
    extreme = G[:, :2, :] @ torch.tensor([(-1.0, -1, 1), (-1, 1, 1), (1, -1, 1), (1, 1, 1)]).t()

    size = torch.tensor((width, height))

    pad_low = ((extreme.min(-1).values + 1) * size).clamp(max=0).abs().ceil().max(0).values.to(torch.int64).tolist()
    pad_high = (extreme.max(-1).values * size - size).clamp(min=0).ceil().max(0).values.to(torch.int64).tolist()

    return pad_low[0], pad_high[0], pad_low[1], pad_high[1]


def try_sample_affine_and_pad(img, p, pad_k, G=None):
    batch, _, height, width = img.shape

    G_try = G

    while True:
        if G is None:
            G_try = sample_affine(p, batch, height, width)

        pad_x1, pad_x2, pad_y1, pad_y2 = get_padding(torch.inverse(G_try), height, width)

        try:
            img_pad = F.pad(img, (pad_x1 + pad_k, pad_x2 + pad_k, pad_y1 + pad_k, pad_y2 + pad_k), mode="reflect",)

        except RuntimeError:
            continue

        break

    return img_pad, G_try, (pad_x1, pad_x2, pad_y1, pad_y2)


def random_apply_affine(img, p, G=None, antialiasing_kernel=SYM6):
    kernel = antialiasing_kernel
    len_k = len(kernel)
    pad_k = (len_k + 1) // 2

    kernel = torch.as_tensor(kernel)
    kernel = torch.ger(kernel, kernel).to(img)
    kernel_flip = torch.flip(kernel, (0, 1))

    img_pad, G, (pad_x1, pad_x2, pad_y1, pad_y2) = try_sample_affine_and_pad(img, p, pad_k, G)

    p_ux1 = pad_x1
    p_ux2 = pad_x2 + 1
    p_uy1 = pad_y1
    p_uy2 = pad_y2 + 1
    w_p = img_pad.shape[3] - len_k + 1
    h_p = img_pad.shape[2] - len_k + 1
    h_o = img.shape[2]
    w_o = img.shape[3]

    img_2x = upfirdn2d(img_pad, kernel_flip, up=2)

    grid = make_grid(
        img_2x.shape,
        -2 * p_ux1 / w_o - 1,
        2 * (w_p - p_ux1) / w_o - 1,
        -2 * p_uy1 / h_o - 1,
        2 * (h_p - p_uy1) / h_o - 1,
        device=img_2x.device,
    ).to(img_2x)
    grid = affine_grid(grid, torch.inverse(G)[:, :2, :].to(img_2x))
    grid = grid * torch.tensor([w_o / w_p, h_o / h_p], device=grid.device) + torch.tensor(
        [(w_o + 2 * p_ux1) / w_p - 1, (h_o + 2 * p_uy1) / h_p - 1], device=grid.device
    )

    img_affine = F.grid_sample(img_2x, grid, mode="bilinear", align_corners=False, padding_mode="zeros")

    img_down = upfirdn2d(img_affine, kernel, down=2)

    end_y = -pad_y2 - 1
    if end_y == 0:
        end_y = img_down.shape[2]

    end_x = -pad_x2 - 1
    if end_x == 0:
        end_x = img_down.shape[3]

    img = img_down[:, :, pad_y1:end_y, pad_x1:end_x]

    return img, G


def apply_color(img, mat):
    batch = img.shape[0]
    img = img.permute(0, 2, 3, 1)
    mat_mul = mat[:, :3, :3].transpose(1, 2).view(batch, 1, 3, 3)
    mat_add = mat[:, :3, 3].view(batch, 1, 1, 3)
    img = img @ mat_mul + mat_add
    img = img.permute(0, 3, 1, 2)

    return img


def random_apply_color(img, p, C=None):
    if C is None:
        C = sample_color(p, img.shape[0])

    img = apply_color(img, C.to(img))

    return img, C


def augment(img, p, transform_matrix=(None, None)):
    img, G = random_apply_affine(img, p, transform_matrix[0])
    img, C = random_apply_color(img, p, transform_matrix[1])

    return img, (G, C)


================================================
FILE: contrastive_learner.py
================================================
import copy
import random
from functools import wraps

import torch
from torch import nn
import torch.nn.functional as F


def identity(x):
    return x


def default(val, def_val):
    return def_val if val is None else val


def flatten(t):
    return t.reshape(t.shape[0], -1)


def safe_concat(arr, el, dim=0):
    if arr is None:
        return el
    return torch.cat((arr, el), dim=dim)


def singleton(cache_key):
    def inner_fn(fn):
        @wraps(fn)
        def wrapper(self, *args, **kwargs):
            instance = getattr(self, cache_key)
            if instance is not None:
                return instance

            instance = fn(self, *args, **kwargs)
            setattr(self, cache_key, instance)
            return instance

        return wrapper

    return inner_fn


# losses


def contrastive_loss(queries, keys, temperature=0.1):
    b, device = queries.shape[0], queries.device
    logits = queries @ keys.t()
    logits = logits - logits.max(dim=-1, keepdim=True).values
    logits /= temperature
    return F.cross_entropy(logits, torch.arange(b, device=device))


def nt_xent_loss(queries, keys, temperature=0.1):
    b, device = queries.shape[0], queries.device

    n = b * 2
    projs = torch.cat((queries, keys))
    logits = projs @ projs.t()

    mask = torch.eye(n, device=device).bool()
    logits = logits[~mask].reshape(n, n - 1)
    logits /= temperature

    labels = torch.cat(((torch.arange(b, device=device) + b - 1), torch.arange(b, device=device)), dim=0)
    loss = F.cross_entropy(logits, labels, reduction="sum")
    loss /= 2 * (b - 1)
    return loss


# augmentation utils


class RandomApply(nn.Module):
    def __init__(self, fn, p):
        super().__init__()
        self.fn = fn
        self.p = p

    def forward(self, x):
        x_out = []
        for ex in x:
            if random.random() > self.p:
                x_out.append(ex[None, :])
            else:
                x_out.append(self.fn(ex))
        return torch.cat(x_out)


# exponential moving average


class EMA:
    def __init__(self, beta):
        super().__init__()
        self.beta = beta

    def update_average(self, old, new):
        if old is None:
            return new
        return old * self.beta + (1 - self.beta) * new


def update_moving_average(ema_updater, ma_model, current_model):
    for current_params, ma_params in zip(current_model.parameters(), ma_model.parameters()):
        old_weight, up_weight = ma_params.data, current_params.data
        ma_params.data = ema_updater.update_average(old_weight, up_weight)


# hidden layer extractor class


class OutputHiddenLayer(nn.Module):
    def __init__(self, net, layer=-2):
        super().__init__()
        self.net = net
        self.layer = layer

        self.hidden = None
        self._register_hook()

    def _find_layer(self):
        if type(self.layer) == str:
            modules = dict([*self.net.named_modules()])
            return modules.get(self.layer, None)
        elif type(self.layer) == int:
            children = [*self.net.children()]
            return children[self.layer]
        elif type(self.layer) == tuple:
            children = [*self.net.children()]
            grand_children = [*children[self.layer[0]].children()]
            return grand_children[self.layer[1]]
        return None

    def _register_hook(self):
        def hook(_, __, output):
            self.hidden = output

        layer = self._find_layer()
        assert layer is not None, f"hidden layer ({self.layer}) not found"
        handle = layer.register_forward_hook(hook)

    def forward(self, x):
        if self.layer == -1:
            return self.net(x)

        _ = self.net(x)
        hidden = self.hidden
        self.hidden = None
        assert hidden is not None, f"hidden layer {self.layer} never emitted an output"
        return hidden


class ContrastiveLearner(nn.Module):
    def __init__(
        self,
        net,
        image_size,
        hidden_layer=-2,
        project_hidden=True,
        project_dim=128,
        use_nt_xent_loss=False,
        use_bilinear=False,
        use_momentum=False,
        momentum_value=0.999,
        key_encoder=None,
        temperature=0.1,
    ):
        super().__init__()
        self.net = OutputHiddenLayer(net, layer=hidden_layer)

        self.temperature = temperature
        self.use_nt_xent_loss = use_nt_xent_loss

        self.project_hidden = project_hidden
        self.projection = None
        self.project_dim = project_dim

        self.use_bilinear = use_bilinear
        self.bilinear_w = None

        self.use_momentum = use_momentum
        self.ema_updater = EMA(momentum_value)
        self.key_encoder = key_encoder

        # for accumulating queries and keys across calls
        self.queries = None
        self.keys = None

        # send a mock image tensor to instantiate parameters
        init = torch.randn(1, 3, image_size, image_size, device="cuda")
        self.forward(init)

    @singleton("key_encoder")
    def _get_key_encoder(self):
        key_encoder = copy.deepcopy(self.net)
        key_encoder._register_hook()
        return key_encoder

    @singleton("bilinear_w")
    def _get_bilinear(self, hidden):
        _, dim = hidden.shape
        return nn.Parameter(torch.eye(dim, device=device, dtype=dtype)).to(hidden)

    @singleton("projection")
    def _get_projection_fn(self, hidden):
        _, dim = hidden.shape
        return nn.Sequential(
            nn.Linear(dim, dim, bias=False), nn.LeakyReLU(inplace=True), nn.Linear(dim, self.project_dim, bias=False)
        ).to(hidden)

    def reset_moving_average(self):
        assert self.use_momentum, "must be using momentum method for key encoder"
        del self.key_encoder
        self.key_encoder = None

    def update_moving_average(self):
        assert self.key_encoder is not None, "key encoder has not been created yet"
        self.key_encoder = update_moving_average(self.ema_updater, self.key_encoder, self.net)

    def calculate_loss(self):
        assert self.queries is not None and self.keys is not None, "no queries or keys accumulated"
        loss_fn = nt_xent_loss if self.use_nt_xent_loss else contrastive_loss
        loss = loss_fn(self.queries, self.keys, temperature=self.temperature)
        self.queries = self.keys = None
        return loss

    def forward(self, x, aug_x, accumulate=False):
        b, c, h, w, device = *x.shape, x.device

        queries = self.net(aug_x)

        key_encoder = self.net if not self.use_momentum else self._get_key_encoder()
        keys = key_encoder(aug_x)

        if self.use_momentum:
            keys = keys.detach()

        queries, keys = map(flatten, (queries, keys))

        if self.use_bilinear:
            W = self._get_bilinear(keys)
            keys = (W @ keys.t()).t()

        project_fn = self._get_projection_fn(queries) if self.project_hidden else identity
        queries, keys = map(project_fn, (queries, keys))

        self.queries = safe_concat(self.queries, queries)
        self.keys = safe_concat(self.keys, keys)

        return self.calculate_loss() if not accumulate else None


================================================
FILE: convert_weight.py
================================================
import argparse
import math
import os
import pickle
import sys

import numpy as np
import torch
from torchvision import utils

from model import Discriminator, Generator


def convert_modconv(vars, source_name, target_name, flip=False):
    weight = vars[source_name + "/weight"].value().eval()
    mod_weight = vars[source_name + "/mod_weight"].value().eval()
    mod_bias = vars[source_name + "/mod_bias"].value().eval()
    noise = vars[source_name + "/noise_strength"].value().eval()
    bias = vars[source_name + "/bias"].value().eval()

    dic = {
        "conv.weight": np.expand_dims(weight.transpose((3, 2, 0, 1)), 0),
        "conv.modulation.weight": mod_weight.transpose((1, 0)),
        "conv.modulation.bias": mod_bias + 1,
        "noise.weight": np.array([noise]),
        "activate.bias": bias,
    }

    dic_torch = {}

    for k, v in dic.items():
        dic_torch[target_name + "." + k] = torch.from_numpy(v)

    if flip:
        dic_torch[target_name + ".conv.weight"] = torch.flip(dic_torch[target_name + ".conv.weight"], [3, 4])

    return dic_torch


def convert_conv(vars, source_name, target_name, bias=True, start=0):
    weight = vars[source_name + "/weight"].value().eval()

    dic = {"weight": weight.transpose((3, 2, 0, 1))}

    if bias:
        dic["bias"] = vars[source_name + "/bias"].value().eval()

    dic_torch = {}

    dic_torch[target_name + f".{start}.weight"] = torch.from_numpy(dic["weight"])

    if bias:
        dic_torch[target_name + f".{start + 1}.bias"] = torch.from_numpy(dic["bias"])

    return dic_torch


def convert_torgb(vars, source_name, target_name):
    weight = vars[source_name + "/weight"].value().eval()
    mod_weight = vars[source_name + "/mod_weight"].value().eval()
    mod_bias = vars[source_name + "/mod_bias"].value().eval()
    bias = vars[source_name + "/bias"].value().eval()

    dic = {
        "conv.weight": np.expand_dims(weight.transpose((3, 2, 0, 1)), 0),
        "conv.modulation.weight": mod_weight.transpose((1, 0)),
        "conv.modulation.bias": mod_bias + 1,
        "bias": bias.reshape((1, 3, 1, 1)),
    }

    dic_torch = {}

    for k, v in dic.items():
        dic_torch[target_name + "." + k] = torch.from_numpy(v)

    return dic_torch


def convert_dense(vars, source_name, target_name):
    weight = vars[source_name + "/weight"].value().eval()
    bias = vars[source_name + "/bias"].value().eval()

    dic = {"weight": weight.transpose((1, 0)), "bias": bias}

    dic_torch = {}

    for k, v in dic.items():
        dic_torch[target_name + "." + k] = torch.from_numpy(v)

    return dic_torch


def update(state_dict, new):
    for k, v in new.items():
        if k not in state_dict:
            raise KeyError(k + " is not found")

        if v.shape != state_dict[k].shape:
            raise ValueError(f"Shape mismatch: {v.shape} vs {state_dict[k].shape}")

        state_dict[k] = v


def discriminator_fill_statedict(statedict, vars, size):
    log_size = int(math.log(size, 2))

    update(statedict, convert_conv(vars, f"{size}x{size}/FromRGB", "convs.0"))

    conv_i = 1

    for i in range(log_size - 2, 0, -1):
        reso = 4 * 2 ** i
        update(
            statedict, convert_conv(vars, f"{reso}x{reso}/Conv0", f"convs.{conv_i}.conv1"),
        )
        update(
            statedict, convert_conv(vars, f"{reso}x{reso}/Conv1_down", f"convs.{conv_i}.conv2", start=1),
        )
        update(
            statedict, convert_conv(vars, f"{reso}x{reso}/Skip", f"convs.{conv_i}.skip", start=1, bias=False),
        )
        conv_i += 1

    update(statedict, convert_conv(vars, f"4x4/Conv", "final_conv"))
    update(statedict, convert_dense(vars, f"4x4/Dense0", "final_linear.0"))
    update(statedict, convert_dense(vars, f"Output", "final_linear.1"))

    return statedict


def fill_statedict(state_dict, vars, size):
    log_size = int(math.log(size, 2))

    for i in range(8):
        update(state_dict, convert_dense(vars, f"G_mapping/Dense{i}", f"style.{i + 1}"))

    update(
        state_dict, {"input.input": torch.from_numpy(vars["G_synthesis/4x4/Const/const"].value().eval())},
    )

    update(state_dict, convert_torgb(vars, "G_synthesis/4x4/ToRGB", "to_rgb1"))

    for i in range(log_size - 2):
        reso = 4 * 2 ** (i + 1)
        update(
            state_dict, convert_torgb(vars, f"G_synthesis/{reso}x{reso}/ToRGB", f"to_rgbs.{i}"),
        )

    update(state_dict, convert_modconv(vars, "G_synthesis/4x4/Conv", "conv1"))

    conv_i = 0

    for i in range(log_size - 2):
        reso = 4 * 2 ** (i + 1)
        update(
            state_dict, convert_modconv(vars, f"G_synthesis/{reso}x{reso}/Conv0_up", f"convs.{conv_i}", flip=True,),
        )
        update(
            state_dict, convert_modconv(vars, f"G_synthesis/{reso}x{reso}/Conv1", f"convs.{conv_i + 1}"),
        )
        conv_i += 2

    for i in range(0, (log_size - 2) * 2 + 1):
        update(
            state_dict, {f"noises.noise_{i}": torch.from_numpy(vars[f"G_synthesis/noise{i}"].value().eval())},
        )

    return state_dict


if __name__ == "__main__":
    device = "cuda"

    parser = argparse.ArgumentParser()
    parser.add_argument("--repo", type=str, required=True)
    parser.add_argument("--gen", action="store_true")
    parser.add_argument("--disc", action="store_true")
    parser.add_argument("--channel_multiplier", type=int, default=2)
    parser.add_argument("path", metavar="PATH")

    args = parser.parse_args()

    sys.path.append(args.repo)

    import dnnlib
    from dnnlib import tflib

    tflib.init_tf()

    with open(args.path, "rb") as f:
        generator, discriminator, g_ema = pickle.load(f)

    size = g_ema.output_shape[2]

    g = Generator(size, 512, 8, channel_multiplier=args.channel_multiplier)
    state_dict = g.state_dict()
    state_dict = fill_statedict(state_dict, g_ema.vars, size)

    g.load_state_dict(state_dict)

    latent_avg = torch.from_numpy(g_ema.vars["dlatent_avg"].value().eval())

    ckpt = {"g_ema": state_dict, "latent_avg": latent_avg}

    if args.gen:
        g_train = Generator(size, 512, 8, channel_multiplier=args.channel_multiplier)
        g_train_state = g_train.state_dict()
        g_train_state = fill_statedict(g_train_state, generator.vars, size)
        ckpt["g"] = g_train_state

    if args.disc:
        disc = Discriminator(size, channel_multiplier=args.channel_multiplier)
        d_state = disc.state_dict()
        d_state = discriminator_fill_statedict(d_state, discriminator.vars, size)
        ckpt["d"] = d_state

    name = os.path.splitext(os.path.basename(args.path))[0]
    torch.save(ckpt, name + ".pt")

    batch_size = {256: 16, 512: 9, 1024: 4}
    n_sample = batch_size.get(size, 25)

    g = g.to(device)

    z = np.random.RandomState(0).randn(n_sample, 512).astype("float32")

    with torch.no_grad():
        img_pt, _ = g([torch.from_numpy(z).to(device)], truncation=0.5, truncation_latent=latent_avg.to(device),)

    Gs_kwargs = dnnlib.EasyDict()
    Gs_kwargs.randomize_noise = False
    img_tf = g_ema.run(z, None, **Gs_kwargs)
    img_tf = torch.from_numpy(img_tf).to(device)

    img_diff = ((img_pt + 1) / 2).clamp(0.0, 1.0) - ((img_tf.to(device) + 1) / 2).clamp(0.0, 1.0)

    img_concat = torch.cat((img_tf, img_pt, img_diff), dim=0)
    utils.save_image(img_concat, name + ".png", nrow=n_sample, normalize=True, range=(-1, 1))


================================================
FILE: dataset.py
================================================
from io import BytesIO

import lmdb
import numpy as np
from PIL import Image
from PIL import Image
from torch.utils.data import Dataset


class MultiResolutionDataset(Dataset):
    def __init__(self, path, transform, resolution=256):
        self.env = lmdb.open(path, max_readers=32, readonly=True, lock=False, readahead=False, meminit=False,)

        if not self.env:
            raise IOError("Cannot open lmdb dataset", path)

        with self.env.begin(write=False) as txn:
            self.length = int(txn.get("length".encode("utf-8")).decode("utf-8"))

        self.resolution = resolution
        self.transform = transform

    def __len__(self):
        return self.length

    def __getitem__(self, index):
        while True:
            try:
                with self.env.begin(write=False) as txn:
                    key = f"{self.resolution}-{str(index).zfill(5)}".encode("utf-8")
                    img_bytes = txn.get(key)

                buffer = BytesIO(img_bytes)
                img = Image.open(buffer)
                break
            except:
                print(f"ERROR loading image {index}")
                index = int(np.random.rand() * self.length)
                print(f"Trying again with {index}...")
        img = self.transform(img)

        return img


================================================
FILE: distributed.py
================================================
import pickle

import torch
from torch import distributed as dist


def get_rank():
    if not dist.is_available():
        return 0

    if not dist.is_initialized():
        return 0

    return dist.get_rank()


def synchronize():
    if not dist.is_available():
        return

    if not dist.is_initialized():
        return

    world_size = dist.get_world_size()

    if world_size == 1:
        return

    dist.barrier()


def get_world_size():
    if not dist.is_available():
        return 1

    if not dist.is_initialized():
        return 1

    return dist.get_world_size()


def reduce_sum(tensor):
    if not dist.is_available():
        return tensor

    if not dist.is_initialized():
        return tensor

    tensor = tensor.clone()
    dist.all_reduce(tensor, op=dist.ReduceOp.SUM)

    return tensor


def gather_grad(params):
    world_size = get_world_size()

    if world_size == 1:
        return

    for param in params:
        if param.grad is not None:
            dist.all_reduce(param.grad.data, op=dist.ReduceOp.SUM)
            param.grad.data.div_(world_size)


def all_gather(data):
    world_size = get_world_size()

    if world_size == 1:
        return [data]

    buffer = pickle.dumps(data)
    storage = torch.ByteStorage.from_buffer(buffer)
    tensor = torch.ByteTensor(storage).to("cuda")

    local_size = torch.IntTensor([tensor.numel()]).to("cuda")
    size_list = [torch.IntTensor([0]).to("cuda") for _ in range(world_size)]
    dist.all_gather(size_list, local_size)
    size_list = [int(size.item()) for size in size_list]
    max_size = max(size_list)

    tensor_list = []
    for _ in size_list:
        tensor_list.append(torch.ByteTensor(size=(max_size,)).to("cuda"))

    if local_size != max_size:
        padding = torch.ByteTensor(size=(max_size - local_size,)).to("cuda")
        tensor = torch.cat((tensor, padding), 0)

    dist.all_gather(tensor_list, tensor)

    data_list = []

    for size, tensor in zip(size_list, tensor_list):
        buffer = tensor.cpu().numpy().tobytes()[:size]
        data_list.append(pickle.loads(buffer))

    return data_list


def reduce_loss_dict(loss_dict):
    world_size = get_world_size()

    if world_size < 2:
        return loss_dict

    with torch.no_grad():
        keys = []
        losses = []

        for k in sorted(loss_dict.keys()):
            keys.append(k)
            losses.append(loss_dict[k])

        losses = torch.stack(losses, 0)
        dist.reduce(losses, dst=0)

        if dist.get_rank() == 0:
            losses /= world_size

        reduced_losses = {k: v for k, v in zip(keys, losses)}

    return reduced_losses


================================================
FILE: generate.py
================================================
import argparse
import torch
from torchvision import utils
from models.stylegan2 import Generator
from tqdm import tqdm


def generate(args, g_ema, device, mean_latent):

    with torch.no_grad():
        g_ema.eval()
        for i in tqdm(range(args.pics)):
            sample_z = torch.randn(args.sample, args.latent, device=device)

            sample, _ = g_ema([sample_z], truncation=args.truncation, truncation_latent=mean_latent)

            utils.save_image(
                sample, f"sample/{str(i).zfill(6)}.png", nrow=1, normalize=True, range=(-1, 1),
            )


if __name__ == "__main__":
    device = "cuda"

    parser = argparse.ArgumentParser()

    parser.add_argument("--size", type=int, default=1024)
    parser.add_argument("--sample", type=int, default=1)
    parser.add_argument("--pics", type=int, default=20)
    parser.add_argument("--truncation", type=float, default=1)
    parser.add_argument("--truncation_mean", type=int, default=4096)
    parser.add_argument("--ckpt", type=str, default="stylegan2-ffhq-config-f.pt")
    parser.add_argument("--channel_multiplier", type=int, default=2)

    args = parser.parse_args()

    args.latent = 512
    args.n_mlp = 8

    g_ema = Generator(args.size, args.latent, args.n_mlp, channel_multiplier=args.channel_multiplier).to(device)
    checkpoint = torch.load(args.ckpt)

    g_ema.load_state_dict(checkpoint["g_ema"])

    if args.truncation < 1:
        with torch.no_grad():
            mean_latent = g_ema.mean_latent(args.truncation_mean)
    else:
        mean_latent = None

    generate(args, g_ema, device, mean_latent)


================================================
FILE: generate_audiovisual.py
================================================
import argparse
import gc
import os
import random
import time
import traceback
import uuid
import warnings

import librosa as rosa
import librosa.display
import numpy as np
import torch as th

import audioreactive as ar
import generate
import render
from models.stylegan1 import G_style
from models.stylegan2 import Generator


def get_noise_range(out_size, generator_resolution, is_stylegan1):
    """Gets the correct number of noise resolutions for a given resolution of StyleGAN 1 or 2"""
    log_max_res = int(np.log2(out_size))
    log_min_res = 2 + (log_max_res - int(np.log2(generator_resolution)))
    if is_stylegan1:
        range_min = log_min_res
        range_max = log_max_res + 1
        side_fn = lambda x: x
    else:
        range_min = 2 * log_min_res + 1
        range_max = 2 * (log_max_res + 1)
        side_fn = lambda x: int(x / 2)
    return range_min, range_max, side_fn


def load_generator(
    ckpt, is_stylegan1, G_res, out_size, noconst, latent_dim, n_mlp, channel_multiplier, dataparallel, base_res_factor
):
    """Loads a StyleGAN 1 or 2 generator"""
    if is_stylegan1:
        generator = G_style(output_size=out_size, checkpoint=ckpt).cuda()
    else:
        generator = Generator(
            G_res,
            latent_dim,
            n_mlp,
            channel_multiplier=channel_multiplier,
            constant_input=not noconst,
            checkpoint=ckpt,
            output_size=out_size,
            base_res_factor=base_res_factor,
        ).cuda()
    if dataparallel:
        generator = th.nn.DataParallel(generator)
    return generator


def generate(
    ckpt,
    audio_file,
    initialize=None,
    get_latents=None,
    get_noise=None,
    get_bends=None,
    get_rewrites=None,
    get_truncation=None,
    output_dir="./output",
    audioreactive_file="audioreactive/examples/default.py",
    offset=0,
    duration=-1,
    latent_file=None,
    shuffle_latents=False,
    G_res=1024,
    out_size=1024,
    fps=30,
    latent_count=12,
    batch=8,
    dataparallel=False,
    truncation=1.0,
    stylegan1=False,
    noconst=False,
    latent_dim=512,
    n_mlp=8,
    channel_multiplier=2,
    randomize_noise=False,
    ffmpeg_preset="slow",
    base_res_factor=1,
    output_file=None,
    args=None,
):
    # if args is empty (i.e. generate() called directly instead of through __main__)
    # create args Namespace with all locally available variables
    if args is None:
        kwargs = locals()
        args = argparse.Namespace()
        for k, v in kwargs.items():
            setattr(args, k, v)

    # ensures smoothing is independent of frame rate
    ar.set_SMF(args.fps / 30)

    time_taken = time.time()
    th.set_grad_enabled(False)

    audio, sr, duration = ar.load_audio(audio_file, offset, duration)

    args.audio = audio
    args.sr = sr

    n_frames = int(round(duration * fps))
    args.duration = duration
    args.n_frames = n_frames

    if initialize is not None:
        args = initialize(args)

    # ====================================================================================
    # =========================== generate audiovisual latents ===========================
    # ====================================================================================
    print("\ngenerating latents...")
    if get_latents is None:
        from audioreactive.default import get_latents

    if latent_file is not None:
        latent_selection = ar.load_latents(latent_file)
    else:
        latent_selection = ar.generate_latents(
            args.latent_count, ckpt, G_res, noconst, latent_dim, n_mlp, channel_multiplier
        )
    if shuffle_latents:
        random_indices = random.sample(range(len(latent_selection)), len(latent_selection))
        latent_selection = latent_selection[random_indices]
    np.save("workspace/last-latents.npy", latent_selection.numpy())

    latents = get_latents(selection=latent_selection, args=args).cpu()

    print(f"{list(latents.shape)} amplitude={latents.std()}\n")

    # ====================================================================================
    # ============================ generate audiovisual noise ============================
    # ====================================================================================
    print("generating noise...")
    if get_noise is None:
        from audioreactive.default import get_noise

    noise = []
    range_min, range_max, exponent = get_noise_range(out_size, G_res, stylegan1)
    for scale in range(range_min, range_max):
        h = (2 if out_size == 1080 else 1) * 2 ** exponent(scale)
        w = (2 if out_size == 1920 else 1) * 2 ** exponent(scale)

        noise.append(get_noise(height=h, width=w, scale=scale - range_min, num_scales=range_max - range_min, args=args))

        if noise[-1] is not None:
            print(list(noise[-1].shape), f"amplitude={noise[-1].std()}")
        gc.collect()
        th.cuda.empty_cache()
    print()

    # ====================================================================================
    # ================ generate audiovisual network bending manipulations ================
    # ====================================================================================
    if get_bends is not None:
        print("generating network bends...")
        bends = get_bends(args=args)
    else:
        bends = []

    # ====================================================================================
    # ================ generate audiovisual model rewriting manipulations ================
    # ====================================================================================
    if get_rewrites is not None:
        print("generating model rewrites...")
        rewrites = get_rewrites(args=args)
    else:
        rewrites = {}

    # ====================================================================================
    # ========================== generate audiovisual truncation =========================
    # ====================================================================================
    if get_truncation is not None:
        print("generating truncation...")
        truncation = get_truncation(args=args)
    else:
        truncation = float(truncation)

    # ====================================================================================
    # ==== render the given (latent, noise, bends, rewrites, truncation) interpolation ===
    # ====================================================================================
    gc.collect()
    th.cuda.empty_cache()

    generator = load_generator(
        ckpt=ckpt,
        is_stylegan1=stylegan1,
        G_res=G_res,
        out_size=out_size,
        noconst=noconst,
        latent_dim=latent_dim,
        n_mlp=n_mlp,
        channel_multiplier=channel_multiplier,
        dataparallel=dataparallel,
        base_res_factor=base_res_factor,
    )

    print(f"\npreprocessing took {time.time() - time_taken:.2f}s\n")

    print(f"rendering {n_frames} frames...")
    if output_file is None:
        checkpoint_title = ckpt.split("/")[-1].split(".")[0].lower()
        track_title = audio_file.split("/")[-1].split(".")[0].lower()
        output_file = f"{output_dir}/{track_title}_{checkpoint_title}_{uuid.uuid4().hex[:8]}.mp4"
    render.render(
        generator=generator,
        latents=latents,
        noise=noise,
        audio_file=audio_file,
        offset=offset,
        duration=duration,
        batch_size=batch,
        truncation=truncation,
        bends=bends,
        rewrites=rewrites,
        out_size=out_size,
        output_file=output_file,
        randomize_noise=randomize_noise,
        ffmpeg_preset=ffmpeg_preset,
    )

    print(f"\ntotal time taken: {(time.time() - time_taken)/60:.2f} minutes")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--ckpt", type=str)
    parser.add_argument("--audio_file", type=str)
    parser.add_argument("--audioreactive_file", type=str, default="audioreactive/examples/default.py")
    parser.add_argument("--output_dir", type=str, default="./output")
    parser.add_argument("--offset", type=float, default=0)
    parser.add_argument("--duration", type=float, default=-1, help="length of rendered video in seconds")
    parser.add_argument("--latent_file", type=str, default=None)
    parser.add_argument("--shuffle_latents", action="store_true")
    parser.add_argument("--G_res", type=int, default=1024)
    parser.add_argument("--out_size", type=int, default=1024, help="rendered video size. Options: 512, 1024, 1920")
    parser.add_argument("--fps", type=int, default=30)
    parser.add_argument("--latent_count", type=int, default=12)
    parser.add_argument("--batch", type=int, default=8)
    parser.add_argument("--dataparallel", action="store_true")
    parser.add_argument("--truncation", type=float, default=1.0)
    parser.add_argument("--stylegan1", action="store_true")
    parser.add_argument("--noconst", action="store_true")
    parser.add_argument("--latent_dim", type=int, default=512)
    parser.add_argument("--n_mlp", type=int, default=8)
    parser.add_argument("--channel_multiplier", type=int, default=2)
    parser.add_argument("--randomize_noise", action="store_true")
    parser.add_argument("--base_res_factor", type=float, default=1)
    parser.add_argument("--ffmpeg_preset", type=str, default="slow")
    parser.add_argument("--output_file", type=str, default=None)
    args = parser.parse_args()

    # ensure output_dir exists
    os.makedirs(args.output_dir, exist_ok=True)

    # transform file path to python module string
    modnames = args.audioreactive_file.replace(".py", "").replace("/", ".").split(".")

    # try to load each of the standard functions from the specified file
    func_names = ["initialize", "get_latents", "get_noise", "get_bends", "get_rewrites", "get_truncation"]
    funcs = {}
    for func in func_names:
        try:
            file = __import__(".".join(modnames[:-1]), fromlist=[modnames[-1]]).__dict__[modnames[-1]]
            funcs[func] = getattr(file, func)
        except AttributeError as error:
            print(f"No '{func}' function found in --audioreactive_file, using default...")
            funcs[func] = None
        except:
            if funcs.get(func, "error") == "error":
                print("Error while loading --audioreactive_file...")
                traceback.print_exc()
                exit(1)

    # override with args from the OVERRIDE dict in the specified file
    arg_dict = vars(args).copy()
    try:
        file = __import__(".".join(modnames[:-1]), fromlist=[modnames[-1]]).__dict__[modnames[-1]]
        for arg, val in getattr(file, "OVERRIDE").items():
            arg_dict[arg] = val
            setattr(args, arg, val)
    except:
        pass  # no overrides, just continue

    ckpt = arg_dict.pop("ckpt", None)
    audio_file = arg_dict.pop("audio_file", None)

    # splat kwargs to function call
    # (generate() has all kwarg defaults specified again to make it amenable to ipynb usage)
    generate(ckpt=ckpt, audio_file=audio_file, **funcs, **arg_dict, args=args)


================================================
FILE: generate_video.py
================================================
import argparse
import uuid

import numpy as np
import torch as th
import torch.multiprocessing as mp
import torch.nn.functional as F

from models.stylegan1 import G_style
from models.stylegan2 import Generator
from render import render


def gaussian_filter(x, sigma):
    dim = len(x.shape)
    if dim != 3 and dim != 4:
        raise Exception("Only 3- or 4-dimensional tensors are supported.")

    radius = sigma * 4
    channels = x.shape[1]

    kernel = th.arange(-radius, radius + 1, dtype=th.float32, device="cuda")
    kernel = th.exp(-0.5 / sigma ** 2 * kernel ** 2)
    kernel = kernel / kernel.sum()
    kernel = kernel.view(1, 1, len(kernel)).repeat(channels, 1, 1)

    if dim == 4:
        t, c, h, w = x.shape
        x = x.view(t, c, h * w)
    x = x.transpose(0, 2)

    x = F.pad(x, (radius, radius), mode="circular")
    x = F.conv1d(x, weight=kernel, groups=channels)

    x = x.transpose(0, 2)
    if dim == 4:
        x = x.view(t, c, h, w)

    return x


def slerp(val, low, high):
    omega = np.arccos(np.clip(np.dot(low / np.linalg.norm(low), high / np.linalg.norm(high)), -1, 1))
    so = np.sin(omega)
    if so == 0:
        return (1.0 - val) * low + val * high  # L'Hopital's rule/LERP
    return np.sin((1.0 - val) * omega) / so * low + np.sin(val * omega) / so * high


def lerp(val, low, high):
    return (1 - val) * low + val * high


def interpolant(t):
    return t * t * t * (t * (t * 6 - 15) + 10)


def perlin_noise(shape, res, tileable=(True, False, False), interpolant=interpolant):
    """Generate a 3D tensor of perlin noise.
    Args:
        shape: The shape of the generated tensor (tuple of three ints).
            This must be a multiple of res.
        res: The number of periods of noise to generate along each
            axis (tuple of three ints). Note shape must be a multiple
            of res.
        tileable: If the noise should be tileable along each axis
            (tuple of three bools). Defaults to (False, False, False).
        interpolant: The interpolation function, defaults to
            t*t*t*(t*(t*6 - 15) + 10).
    Returns:
        A tensor of shape shape with the generated noise.
    Raises:
        ValueError: If shape is not a multiple of res.
    """
    delta = (res[0] / shape[0], res[1] / shape[1], res[2] / shape[2])
    d = (shape[0] // res[0], shape[1] // res[1], shape[2] // res[2])
    grid = np.mgrid[0 : res[0] : delta[0], 0 : res[1] : delta[1], 0 : res[2] : delta[2]]
    # print(np.mgrid[0 : res[0] : delta[0]])
    # print(0, res[0], delta[0])
    # print(th.linspace(0, res[0], delta[0]))
    # grid = th.meshgrid(
    #     th.linspace(0, res[0], delta[0]), th.linspace(0, res[1], delta[1]), th.linspace(0, res[1], delta[1])
    # ).cuda()
    grid = grid.transpose(1, 2, 3, 0) % 1
    grid = th.from_numpy(grid).cuda()
    # Gradients
    theta = 2 * np.pi * np.random.rand(res[0] + 1, res[1] + 1, res[2] + 1)
    phi = 2 * np.pi * np.random.rand(res[0] + 1, res[1] + 1, res[2] + 1)
    gradients = np.stack((np.sin(phi) * np.cos(theta), np.sin(phi) * np.sin(theta), np.cos(phi)), axis=3)
    if tileable[0]:
        gradients[-1, :, :] = gradients[0, :, :]
    if tileable[1]:
        gradients[:, -1, :] = gradients[:, 0, :]
    if tileable[2]:
        gradients[:, :, -1] = gradients[:, :, 0]
    gradients = gradients.repeat(d[0], 0).repeat(d[1], 1).repeat(d[2], 2)
    gradients = th.from_numpy(gradients).cuda()
    g000 = gradients[: -d[0], : -d[1], : -d[2]]
    g100 = gradients[d[0] :, : -d[1], : -d[2]]
    g010 = gradients[: -d[0], d[1] :, : -d[2]]
    g110 = gradients[d[0] :, d[1] :, : -d[2]]
    g001 = gradients[: -d[0], : -d[1], d[2] :]
    g101 = gradients[d[0] :, : -d[1], d[2] :]
    g011 = gradients[: -d[0], d[1] :, d[2] :]
    g111 = gradients[d[0] :, d[1] :, d[2] :]
    # Ramps
    n000 = th.sum(th.stack((grid[:, :, :, 0], grid[:, :, :, 1], grid[:, :, :, 2]), axis=3) * g000, 3)
    n100 = th.sum(th.stack((grid[:, :, :, 0] - 1, grid[:, :, :, 1], grid[:, :, :, 2]), axis=3) * g100, 3)
    n010 = th.sum(th.stack((grid[:, :, :, 0], grid[:, :, :, 1] - 1, grid[:, :, :, 2]), axis=3) * g010, 3)
    n110 = th.sum(th.stack((grid[:, :, :, 0] - 1, grid[:, :, :, 1] - 1, grid[:, :, :, 2]), axis=3) * g110, 3)
    n001 = th.sum(th.stack((grid[:, :, :, 0], grid[:, :, :, 1], grid[:, :, :, 2] - 1), axis=3) * g001, 3)
    n101 = th.sum(th.stack((grid[:, :, :, 0] - 1, grid[:, :, :, 1], grid[:, :, :, 2] - 1), axis=3) * g101, 3)
    n011 = th.sum(th.stack((grid[:, :, :, 0], grid[:, :, :, 1] - 1, grid[:, :, :, 2] - 1), axis=3) * g011, 3)
    n111 = th.sum(th.stack((grid[:, :, :, 0] - 1, grid[:, :, :, 1] - 1, grid[:, :, :, 2] - 1), axis=3) * g111, 3)
    # Interpolation
    t = interpolant(grid)
    n00 = n000 * (1 - t[:, :, :, 0]) + t[:, :, :, 0] * n100
    n10 = n010 * (1 - t[:, :, :, 0]) + t[:, :, :, 0] * n110
    n01 = n001 * (1 - t[:, :, :, 0]) + t[:, :, :, 0] * n101
    n11 = n011 * (1 - t[:, :, :, 0]) + t[:, :, :, 0] * n111
    n0 = (1 - t[:, :, :, 1]) * n00 + t[:, :, :, 1] * n10
    n1 = (1 - t[:, :, :, 1]) * n01 + t[:, :, :, 1] * n11
    return (1 - t[:, :, :, 2]) * n0 + t[:, :, :, 2] * n1


def spline_loops(base_latent_selection, loop_starting_latents, n_frames, num_loops, smoothing, s=True):
    from scipy import interpolate

    base_latent_selection = np.concatenate([base_latent_selection, base_latent_selection[[0]]])

    x = np.linspace(0, 1, n_frames // max(1, num_loops))
    base_latents = np.zeros((len(x), *base_latent_selection.shape[1:]))
    for lay in range(base_latent_selection.shape[1]):
        for lat in range(base_latent_selection.shape[2]):
            tck = interpolate.splrep(
                np.linspace(0, 1, base_latent_selection.shape[0], dtype=np.float32), base_latent_selection[:, lay, lat],
            )
            base_latents[:, lay, lat] = interpolate.splev(x, tck)

    base_latents = th.cat([th.from_numpy(base_latents)] * int(n_frames / len(base_latents)), axis=0).float()

    return base_latents


def get_latent_loops(base_latent_selection, loop_starting_latents, n_frames, num_loops, smoothing, s=True):
    base_latents = []

    for n in range(len(base_latent_selection)):
        for val in np.linspace(0.0, 1.0, int(n_frames // max(1, num_loops) // len(base_latent_selection))):
            base_latents.append(
                (slerp if s else lerp)(
                    val,
                    base_latent_selection[(n + loop_starting_latents) % len(base_latent_selection)][0].cpu(),
                    base_latent_selection[(n + loop_starting_latents + 1) % len(base_latent_selection)][0].cpu(),
                )
            )

    base_latents = th.stack(base_latents, axis=0).cuda()
    base_latents = th.cat([base_latents] * int(n_frames / len(base_latents)), axis=0)
    base_latents = th.stack([base_latents] * base_latent_selection.shape[1], axis=1)

    base_latents = gaussian_filter(base_latents, smoothing)

    return base_latents


if "main" in __name__:
    parser = argparse.ArgumentParser()

    parser.add_argument("--ckpt", type=str)
    parser.add_argument("--G_res", type=int, default=1024)
    parser.add_argument("--out_size", type=int, default=1024)
    parser.add_argument("--batch", type=int, default=12)
    parser.add_argument("--n_frames", type=int, default=24 * 30)
    parser.add_argument("--duration", type=int, default=24)
    parser.add_argument("--const", type=bool, default=False)
    parser.add_argument("--channel_multiplier", type=int, default=2)
    parser.add_argument("--truncation", type=int, default=0.7)
    parser.add_argument("--stylegan1", type=bool, default=False)
    parser.add_argument("--slerp", type=bool, default=True)
    parser.add_argument("--latents", type=str, default=None)

    args = parser.parse_args()

    th.set_grad_enabled(False)
    th.backends.cudnn.benchmark = True
    mp.set_start_method("spawn")

    if args.stylegan1:
        generator = G_style(output_size=args.out_size, checkpoint=args.ckpt).cuda()
    else:
        args.latent = 512
        args.n_mlp = 8
        generator = Generator(
            args.G_res,
            args.latent,
            args.n_mlp,
            channel_multiplier=args.channel_multiplier,
            constant_input=args.const,
            checkpoint=args.ckpt,
            output_size=args.out_size,
        ).cuda()
    # generator = th.nn.DataParallel(generator)

    if args.latents is not None:
        styles = th.from_numpy(np.load(args.latents))
    else:
        # styles1 = th.randn((int(args.duration / 3), 512), device="cuda")
        # styles1 = generator(styles1, map_latents=True)
        # styles2 = th.randn((int(args.duration / 3), 512), device="cuda")
        # styles2 = generator(styles2, map_latents=True)
        # styles3 = th.randn((int(args.duration / 3), 512), device="cuda")
        # styles3 = generator(styles3, map_latents=True)

        styles = th.randn((args.duration, 512), device="cuda")
        styles = generator(styles, map_latents=True)

    latents = th.cat([styles[[0]]] * args.n_frames, axis=0)

    # moving_low = spline_loops(
    #     styles1.cpu(), 0, int(args.n_frames / 3), num_loops=1, smoothing=1, s=args.slerp
    # ).cuda()[:, :5]
    # moving_mid = spline_loops(
    #     styles2.cpu(), 0, int(args.n_frames / 3), num_loops=1, smoothing=1, s=args.slerp
    # ).cuda()[:, 5:10]
    # moving_hi = spline_loops(
    #     styles3.cpu(), 0, int(args.n_frames / 3), num_loops=1, smoothing=1, s=args.slerp
    # ).cuda()[:, 10:]
    # static_low = th.cat([moving_low[[0]]] * int(args.n_frames / 3), axis=0)
    # static_mid = th.cat([moving_mid[[0]]] * int(args.n_frames / 3), axis=0)
    # static_hi = th.cat([moving_hi[[0]]] * int(args.n_frames / 3), axis=0)

    # print(
    #     th.cat([moving_low, static_mid, static_hi], axis=1).shape,
    #     th.cat([static_low[:60], static_mid[:60], static_hi[:60]], axis=1).shape,
    #     th.cat([static_low, moving_mid, static_hi], axis=1).shape,
    #     th.cat([static_low[:60], static_mid[:60], static_hi[:60]], axis=1).shape,
    #     th.cat([static_low, static_mid, moving_hi], axis=1).shape,
    # )

    # print(th.cat([static_low[[0]], static_mid[[0]], static_hi[[0]]], axis=1).cpu().numpy().shape)
    # np.save("latents_example.npy", th.cat([static_low[[0]], static_mid[[0]], static_hi[[0]]], axis=1).cpu().numpy())

    # latents = th.cat(
    #     [
    #         th.cat([moving_low, static_mid, static_hi], axis=1),
    #         th.cat([static_low[:60], static_mid[:60], static_hi[:60]], axis=1),
    #         th.cat([static_low, moving_mid, static_hi], axis=1),
    #         th.cat([static_low[:60], static_mid[:60], static_hi[:60]], axis=1),
    #         th.cat([static_low, static_mid, moving_hi], axis=1),
    #     ],
    #     axis=0,
    # ).float()

    # latents = gaussian_filter(latents, 7)

    latents = latents.cpu()

    print("latent shape: ")
    print(latents.shape, "\n")

    log_max_res = int(np.log2(args.out_size))
    log_min_res = 2 + (log_max_res - int(np.log2(args.G_res)))

    noise = []
    if args.stylegan1:
        for s in range(log_min_res, log_max_res + 1):
            h = 2 ** s
            w = (2 if args.out_size == 1920 else 1) * 2 ** s
            noise.append(th.randn((1, 1, h, w), device="cuda"))
    else:
        for s in range(2 * log_min_res + 1, 2 * (log_max_res + 1), 1):
            h = 2 ** int(s / 2)
            w = (2 if args.out_size == 1920 else 1) * 2 ** int(s / 2)
            noise.append(th.randn((1, 1, h, w), device="cuda"))

    def create_circular_mask(h, w, center=None, radius=None):
        if center is None:  # use the middle of the image
            center = (int(w / 2), int(h / 2))
        if radius is None:  # use the smallest distance between the center and image walls
            radius = min(center[0], center[1], w - center[0], h - center[1])
        Y, X = np.ogrid[:h, :w]
        dist_from_center = np.sqrt((X - center[0]) ** 2 + (Y - center[1]) ** 2)
        mask = dist_from_center <= radius
        return th.from_numpy(mask)

    print("noise shapes: ")
    for i, n in enumerate(noise):
        if n is None:
            continue
        if i > 14:
            noise[i] = None
            continue

        # mask = create_circular_mask(n.shape[-2], n.shape[-1], radius=n.shape[-1] / 2.5)[None, ...].float()
        # mask = th.stack(
        #     [
        #         th.cat(
        #             [
        #                 th.zeros((int(n.shape[-2] * 1 / 2))),
        #                 th.linspace(0, 1, int(n.shape[-2] * 1 / 4)),
        #                 th.ones((int(n.shape[-2] * 1 / 4))),
        #             ],
        #             axis=0,
        #         )
        #     ]
        #     * n.shape[-1]
        # ).T[None, ...]
        # mask = th.stack([mask] * n.shape[0], axis=0)
        # noise[i] = mask * n[[0]].cpu()  # gaussian_filter(n, 24).cpu()

        if i < 4:
            moving = 2 * gaussian_filter(th.randn((200, 1, n.shape[-2], n.shape[-1]), device="cuda"), 3)
            # moving /= moving.std()
            static = th.cat([n] * (len(latents) - len(moving)))
            print(moving.shape, static.shape)
            # static /= static.std()
            noise[i] = th.cat([moving, static], axis=0)
        elif 4 <= i < 8:
            static1 = th.cat([n] * (260))
            # static1 /= static1.std()
            moving = 4 * gaussian_filter(th.randn((200, 1, n.shape[-2], n.shape[-1]), device="cuda"), 3)
            # moving /= moving.std()
            static2 = th.cat([n] * (len(latents) - 460))
            print(static1.shape, moving.shape, static2.shape)
            # static2 /= static2.std()
            noise[i] = th.cat([static1, moving, static2], axis=0)
        elif i >= 8:
            moving = 8 * gaussian_filter(th.randn((200, 1, n.shape[-2], n.shape[-1]), device="cuda"), 3)
            # moving /= moving.std()
            static = th.cat([n] * (len(latents) - len(moving)))
            print(static.shape, moving.shape)
            # static /= static.std()
            noise[i] = th.cat([static, moving], axis=0)
        noise[i] = gaussian_filter(noise[i].cuda(), 7).cpu()

        # noise[i] = th.cat([n[[0]]] * len(latents), axis=0).cpu()  # gaussian_filter(n, 24).cpu()
        # noise[i] /= noise[i].std()

        # if i > 2 and i < 13:
        #     # xs = 8 * np.pi * th.linspace(0, 1, n.shape[-1])
        #     # ys = th.linspace(0, 2 * np.pi, n.shape[-2])
        #     # ts = 8 * np.pi * th.linspace(0, 1, n.shape[0])
        #     # horiz = xs[None, None, None, :] + ys[None, None, :, None] + ts[:, None, None, None]
        #     # vert = (
        #     #     xs[None, None, None, :] / (4 * np.pi)
        #     #     + 4 * np.pi * ys[None, None, :, None]
        #     #     + 2 * ts[:, None, None, None]
        #     # )
        #     # moving_noise = th.sin(horiz.cuda() * vert.cuda() + n / 4)
        #     # moving_noise = gaussian_filter(moving_noise, 6).cpu()
        #     # moving_noise /= moving_noise.std() / 2
        #     moving_noise = perlin_noise((n.shape[0], n.shape[-2], n.shape[-1]), (10, 8, 8))[:, None, ...]
        #     moving_noise += gaussian_filter(n, 8) / 2.5
        #     moving_noise /= moving_noise.std() / 1.5
        #     noise[i] += (1 - mask) * moving_noise.cpu()

        print(i, noise[i].shape, noise[i].std())
    print()

    import ffmpeg

    output_name = f"/home/hans/neurout/{args.ckpt.split('/')[-1].split('.')[0]}_{uuid.uuid4().hex[:8]}"

    video = (
        ffmpeg.input("pipe:", format="rawvideo", pix_fmt="rgb24", framerate=len(latents) / args.duration, s="256x256")
        .output(f"{output_name}_noise.mp4", framerate=len(latents) / args.duration, vcodec="libx264", preset="slow",)
        .global_args("-benchmark", "-stats", "-hide_banner")
        .overwrite_output()
        .run_async(pipe_stdin=True)
    )
    print(
        noise[3][:200].shape,
        noise[3][-45:-15].shape,
        noise[7][15:45].shape,
        noise[7][260:460].shape,
        noise[7][15:45].shape,
        noise[12][15:45].shape,
        noise[12][520:].shape,
    )
    # output = noise[-5].permute(0, 2, 3, 1).numpy()
    output = th.cat(
        [
            F.interpolate(noise[3][:200], (256, 256)),
            F.interpolate(th.cat([noise[3][[200]]] * 30, axis=0), (256, 256)),
            F.interpolate(th.cat([noise[7][[260]]] * 30, axis=0), (256, 256)),
            F.interpolate(noise[7][260:460], (256, 256)),
            F.interpolate(th.cat([noise[7][[460]]] * 30, axis=0), (256, 256)),
            F.interpolate(th.cat([noise[12][[520]]] * 30, axis=0), (256, 256)),
            F.interpolate(noise[12][520:], (256, 256)),
        ],
        axis=0,
    )
    print(output.shape)
    output = output.permute(0, 2, 3, 1).numpy()
    print(output.shape)
    output = output / output.max()
    output = output - output.min()
    output = output * 255
    output = output.astype(np.uint8)
    output = np.concatenate([output] * 3, axis=3)
    for frame in output:
        video.stdin.write(frame.tobytes())
    video.stdin.close()
    video.wait()
    # video.close()

    # video = (
    #     ffmpeg.input("pipe:", format="rawvideo", pix_fmt="rgb24", framerate=len(latents) / args.duration, s="23x23")
    #     .output(f"{output_name}_latents.mp4", framerate=len(latents) / args.duration, vcodec="libx264", preset="slow",)
    #     .global_args("-benchmark", "-stats", "-hide_banner")
    #     .overwrite_output()
    #     .run_async(pipe_stdin=True)
    # )
    # output = th.cat(
    #     [
    #         latents[: int(args.n_frames / 3), 0],  #                                lo
    #         latents[-45:-15, 0],  #                                                      pause lo
    #         latents[15:45, 7],  #                                                       pause mid
    #         latents[60 + int(args.n_frames / 3) : 60 + int(2 * args.n_frames / 3), 7],  #   mid
    #         latents[15:45, 7],  #                                                       pause mid
    #         latents[15:45, 14],  #                                                      pause hit
    #         latents[120 + int(2 * args.n_frames / 3) :, 14],  #                           hi
    #     ],
    #     axis=0,
    # )
    # print(output.shape)
    # output = th.cat([output, th.zeros((len(latents), 17))], axis=1)
    # print(output.shape)
    # output = output.reshape((len(latents), 23, 23, 1)).numpy()
    # print(output.shape)
    # output = output / output.max()
    # output = output - output.min()
    # output = output * 255
    # output = output.astype(np.uint8)
    # output = np.concatenate([output] * 3, axis=3)
    # for frame in output:
    #     video.stdin.write(frame.tobytes())
    # video.stdin.close()
    # video.wait()

    # noise = []
    # if args.stylegan1:
    #     for s in range(log_min_res, log_max_res + 1):
    #         h = 2 ** s
    #         w = (2 if args.out_size == 1920 else 1) * 2 ** s
    #         noise.append(np.random.normal(size=(args.n_frames, 1, h, w)))
    # else:
    #     for s in range(2 * log_min_res + 1, 2 * (log_max_res + 1), 1):
    #         h = 2 ** int(s / 2)
    #         w = (2 if args.out_size == 1920 else 1) * 2 ** int(s / 2)
    #         noise.append(np.random.normal(size=(args.n_frames, 1, h, w)))

    # print("noise shapes: ")
    # for i, n in enumerate(noise):
    #     if n is None:
    #         continue
    #     noise[i] = th.from_numpy(ndi.gaussian_filter(n, [15, 0, 0, 0], mode="wrap"))
    #     print(n.shape)
    # print()

    class addNoise(th.nn.Module):
        def __init__(self, noise):
            super(addNoise, self).__init__()
            self.noise = noise

        def forward(self, x):
            return x + self.noise

    manipulations = []
    if log_min_res > 2:
        reflects = []
        for lres in range(2, log_min_res):
            half = 2 ** (lres - 1)
            reflects.append(th.nn.ReplicationPad2d((half, half, half, half)))
        manipulations += [
            {
                "layer": 0,
                "transform": th.nn.Sequential(
                    *reflects, addNoise(2 * th.randn(size=(1, 1, 2 ** log_min_res, 2 ** log_min_res), device="cuda")),
                ),
            }
        ]

    # tl = 4
    # width = lambda s: (2 if args.out_size == 1920 else 1) * 2 ** int(s)
    # translation = (
    #     th.tensor([np.linspace(0, width(tl), args.n_frames + 1), np.zeros((args.n_frames + 1,))]).float().T[:-1]
    # )
    # manipulations += [{"layer": tl, "transform": "translateX", "params": translation}]

    # zl = 6
    # print(
    #     th.cat(
    #         [
    #             th.linspace(-1, 3, int(args.n_frames / 2)),
    #             th.linspace(3, -1, args.n_frames - int(args.n_frames / 2)) + 1,
    #         ]
    #     ).shape
    # )
    # zoom = gaussian_filter(
    #     th.cat(
    #         [
    #             th.linspace(0, 3, int(args.n_frames / 2), dtype=th.float32, device="cuda"),
    #             th.linspace(3, 0, args.n_frames - int(args.n_frames / 2), dtype=th.float32, device="cuda") + 1,
    #         ]
    #     )[:, None, None],
    #     30,
    # ).squeeze()
    # zoom -= zoom.min()
    # zoom /= zoom.max()
    # # zoom *= 1.5
    # zoom += 0.5
    # print(zoom.min().item(), zoom.max().item(), zoom.shape)
    # manipulations += [{"layer": zl, "transform": "zoom", "params": zoom}]

    # rl = 6
    # rotation = th.nn.Sigmoid()(th.tensor(np.linspace(0.0, 1.0, args.n_frames + 1), device="cuda").float())
    # rotation -= rotation.min()
    # rotation /= rotation.max()
    # rotation = rotation[:-1]
    # manipulations += [{"layer": rl, "transform": "rotate", "params": (360.0 * rotation).cpu()}]

    render(
        generator=generator,
        latents=latents,
        noise=noise,
        offset=0,
        duration=args.duration,
        batch_size=args.batch,
        truncation=args.truncation,
        manipulations=manipulations,
        out_size=args.out_size,
        output_file=f"{output_name}.mp4",
    )


================================================
FILE: gpu_profile.py
================================================
import datetime
import linecache
import os

os.environ["CUDA_LAUNCH_BLOCKING"] = "1"

from py3nvml import py3nvml
import torch
import socket

# different settings
print_tensor_sizes = True
use_incremental = False

if "GPU_DEBUG" in os.environ:
    gpu_profile_fn = f"host_{socket.gethostname()}_gpu{os.environ['GPU_DEBUG']}_mem_prof-{datetime.datetime.now():%d-%b-%y-%H-%M-%S}.prof.txt"
    print("profiling gpu usage to ", gpu_profile_fn)

## Global variables
last_tensor_sizes = set()
last_meminfo_used = 0
lineno = None
func_name = None
filename = None
module_name = None


def gpu_profile(frame, event, arg):
    # it is _about to_ execute (!)
    global last_tensor_sizes
    global last_meminfo_used
    global lineno, func_name, filename, module_name

    if event == "line":
        try:
            # about _previous_ line (!)
            if lineno is not None:
                py3nvml.nvmlInit()
                handle = py3nvml.nvmlDeviceGetHandleByIndex(int(os.environ["GPU_DEBUG"]))
                meminfo = py3nvml.nvmlDeviceGetMemoryInfo(handle)
                line = linecache.getline(filename, lineno)
                where_str = module_name + " " + func_name + ":" + str(lineno)

                new_meminfo_used = meminfo.used
                mem_display = new_meminfo_used - last_meminfo_used if use_incremental else new_meminfo_used
                if abs(new_meminfo_used - last_meminfo_used) / 1024 ** 2 > 256:
                    with open(gpu_profile_fn, "a+") as f:
                        f.write(f"{where_str:<50}" f":{(mem_display)/1024**2:<7.1f}Mb " f"{line.rstrip()}\n")

                        last_meminfo_used = new_meminfo_used
                        if print_tensor_sizes is True:
                            for tensor in get_tensors():
                                if not hasattr(tensor, "dbg_alloc_where"):
                                    tensor.dbg_alloc_where = where_str
                            new_tensor_sizes = {(type(x), tuple(x.size()), x.dbg_alloc_where) for x in get_tensors()}
                            for t, s, loc in new_tensor_sizes - last_tensor_sizes:
                                f.write(f"+ {loc:<50} {str(s):<20} {str(t):<10}\n")
                            for t, s, loc in last_tensor_sizes - new_tensor_sizes:
                                f.write(f"- {loc:<50} {str(s):<20} {str(t):<10}\n")
                            last_tensor_sizes = new_tensor_sizes
                py3nvml.nvmlShutdown()

            # save details about line _to be_ executed
            lineno = None

            func_name = frame.f_code.co_name
            filename = frame.f_globals["__file__"]
            if filename.endswith(".pyc") or filename.endswith(".pyo"):
                filename = filename[:-1]
            module_name = frame.f_globals["__name__"]
            lineno = frame.f_lineno

            # only profile codes within the parent folder, otherwise there are too many function calls into other pytorch scripts
            # need to modify the key words below to suit your case.
            if "maua-stylegan2" not in os.path.dirname(os.path.abspath(filename)):
                lineno = None  # skip current line evaluation

            if (
                "car_datasets" in filename
                or "_exec_config" in func_name
                or "gpu_profile" in module_name
                or "tee_stdout" in module_name
                or "PIL" in module_name
            ):
                lineno = None  # skip othe unnecessary lines

            return gpu_profile

        except (KeyError, AttributeError):
            pass

    return gpu_profile


def get_tensors(gpu_only=True):
    import gc

    for obj in gc.get_objects():
        try:
            if torch.is_tensor(obj):
                tensor = obj
            elif hasattr(obj, "data") and torch.is_tensor(obj.data):
                tensor = obj.data
            else:
                continue

            if tensor.is_cuda:
                yield tensor
        except Exception as e:
            pass


================================================
FILE: gpumon.py
================================================
import argparse
import os
import signal
import subprocess
import time
from queue import Empty, Queue
from threading import Thread

import numpy as np
import wandb

parser = argparse.ArgumentParser()

parser.add_argument("--wbname", type=str, required=True)
parser.add_argument("--wbproj", type=str, required=True)
parser.add_argument("--wbgroup", type=str, default=None)

args = parser.parse_args()

if args.wbgroup is None:
    wandb.init(project=args.wbproj, name=args.wbname, settings=wandb.Settings(_disable_stats=True))
else:
    wandb.init(project=args.wbproj, group=args.wbgroup, name=args.wbname, settings=wandb.Settings(_disable_stats=True))


def enqueue_output(out, queue):
    for line in iter(out.readline, b""):
        queue.put(line)
    out.close()


os.setpgrp()

clock_proc = subprocess.Popen("nvidia-smi dmon -s c", shell=True, stdout=subprocess.PIPE, bufsize=1)
clock_proc.daemon = True

time.sleep(0.5)

throttle_reasons = [
    "clocks_throttle_reasons.gpu_idle",
    "clocks_throttle_reasons.applications_clocks_setting",
    "clocks_throttle_reasons.sw_power_cap",
    "clocks_throttle_reasons.sw_thermal_slowdown",
    "clocks_throttle_reasons.hw_slowdown",
    "clocks_throttle_reasons.hw_thermal_slowdown",
    "clocks_throttle_reasons.hw_power_brake_slowdown",
    "clocks_throttle_reasons.sync_boost",
]
throttle_proc = subprocess.Popen(
    f"nvidia-smi --query-gpu=index,{','.join(throttle_reasons)} --format=csv,noheader --loop=1",
    shell=True,
    stdout=subprocess.PIPE,
    bufsize=1,
)
throttle_proc.daemon = True

# create queue that gets the output lines from both processes
q = Queue()
clock_thread = Thread(target=enqueue_output, args=(clock_proc.stdout, q))
clock_thread.daemon = True
thottle_thread = Thread(target=enqueue_output, args=(throttle_proc.stdout, q))
thottle_thread.daemon = True

clock_thread.start()
thottle_thread.start()

throttles = [[], []]
clocks = [[], []]
while clock_proc.poll() is None or not q.empty():
    try:
        line = q.get_nowait()
    except Empty:
        pass
    else:
        line = line.decode("utf-8").strip()
        if "#" in line:
            continue
        if "," in line:
            raw = line.split(",")
            gpu = int(raw[0])
            bits = [0 if "Not" in a else 1 for a in raw[1:]]
            throttles[gpu].append(bits)
            # print(gpu, bits)
        else:
            raw = line.split("  ")
            gpu = int(raw[0])
            clock = int(raw[-1])
            clocks[gpu].append(clock)
            # print(gpu, clock)

    if len(clocks[0]) > 30:
        try:
            throttles = np.array(throttles)
            clocks = np.array(clocks)
            log_dict = {}
            for gpu in [0, 1]:
                log_dict[f"gpu.{gpu}.clock.speed"] = np.mean(clocks[gpu])

                for r, reason in enumerate(throttle_reasons):
                    log_dict[f"gpu.{gpu}.{reason}"] = np.mean(throttles[gpu, :, r])

            print("\n".join([k.ljust(80) + str(v) for k, v in log_dict.items()]))
            wandb.log(log_dict)
        except:
            pass

        throttles = [[], []]
        clocks = [[], []]

os.kill(throttle_proc.pid, signal.SIGINT)
os.kill(clock_proc.pid, signal.SIGINT)


================================================
FILE: lightning.py
================================================
import os
import gc
import math
import wandb
import random
import argparse
import validation
import torch as th
import torch.nn as nn
import torch.nn.functional as F
import torchvision as tv
import torchvision.transforms as transforms
from collections import OrderedDict
from torch.utils import data
from dataset import MultiResolutionDataset
from model import Generator, Discriminator
import pytorch_lightning as pl


def requires_grad(model, flag=True):
    for p in model.parameters():
        p.requires_grad = flag


def get_spectral_norms(model):
    spectral_norms = {}
    for name, param in model.named_parameters():
        if param.numel() > 0:
            spectral_norms[name] = nn.utils.spectral_norm(param)
    return spectral_norms


class StyleGAN2(pl.LightningModule):
    def __init__(self, hparams):
        super().__init__()
        self.hparams = hparams  # for automatic param saving with lightning
        [setattr(self, k, v) for k, v in vars(hparams).items()]  # for easy access within module

        self.generator = Generator(self.size, self.latent_size, self.n_mlp, channel_multiplier=self.channel_multiplier)
        self.g_ema = Generator(self.size, self.latent_size, self.n_mlp, channel_multiplier=self.channel_multiplier)
        self.g_ema.eval()
        self.accumulate_g(0)

        self.discriminator = Discriminator(self.size, channel_multiplier=self.channel_multiplier)

        self.sample_z = th.randn(self.n_sample, self.latent_size)

        self.mean_path_length = th.tensor(0.0)

    def forward(self, z):
        return self.generator(z)

    def accumulate_g(self, decay=0.5 ** (32.0 / (10_000))):
        par1 = dict(self.g_ema.named_parameters())
        par2 = dict(self.generator.named_parameters())
        for name, param in self.g_ema.named_parameters():
            param.data = decay * par1[name].data + (1 - decay) * par2[name].data

    def configure_optimizers(self):
        g_reg_ratio = self.g_reg_every / (self.g_reg_every + 1)
        d_reg_ratio = self.d_reg_every / (self.d_reg_every + 1)

        g_optim = th.optim.Adam(
            self.generator.parameters(), lr=self.lr * g_reg_ratio, betas=(0 ** g_reg_ratio, 0.99 ** g_reg_ratio),
        )
        d_optim = th.optim.Adam(
            self.discriminator.parameters(), lr=self.lr * d_reg_ratio, betas=(0 ** d_reg_ratio, 0.99 ** d_reg_ratio),
        )

        return [g_optim, g_optim, d_optim, d_optim], []

    def configure_apex(self, amp, model, optimizers, amp_level):
        amp_optimizers = []
        for optimizer in optimizers:
            try:
                amp_model, amp_optimizer = amp.initialize(model, optimizer, opt_level=amp_level,)
            except RuntimeError as err:
                print(err)
                print("Skipping this optimizer")
            amp_optimizers.append(amp_optimizer)
        return amp_model, amp_optimizers

    def train_dataloader(self):
        transform = transforms.Compose(
            [
                transforms.RandomVerticalFlip(p=0.5 if self.vflip else 0),
                transforms.RandomHorizontalFlip(p=0.5 if self.hflip else 0),
                transforms.ToTensor(),
                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True),
            ]
        )
        dataset = MultiResolutionDataset(self.path, transform, self.size)
        loader = data.DataLoader(dataset, batch_size=self.batch_size, shuffle=True, num_workers=0)
        return loader

    def d_logistic_loss(self, real_pred, fake_pred):
        real_loss = F.softplus(-real_pred)
        fake_loss = F.softplus(fake_pred)
        return real_loss.mean() + fake_loss.mean()

    def d_r1_loss(self, real_pred, real_img):
        (grad_real,) = th.autograd.grad(outputs=real_pred.sum(), inputs=real_img, create_graph=True)
        grad_penalty = grad_real.pow(2).view(grad_real.shape[0], -1).sum(1).mean()
        return grad_penalty

    def g_nonsaturating_loss(self, fake_pred):
        loss = F.softplus(-fake_pred).mean()
        return loss

    def g_path_regularize(self, fake_img, latents, mean_path_length, decay=0.01):
        noise = th.randn_like(fake_img) / math.sqrt(fake_img.shape[2] * fake_img.shape[3])
        # print(fake_img.requires_grad, noise.requires_grad, latents.requires_grad)
        (grad,) = th.autograd.grad(outputs=(fake_img * noise).sum(), inputs=latents, create_graph=True)
        path_lengths = th.sqrt(grad.pow(2).sum(2).mean(1))
        path_mean = mean_path_length + decay * (path_lengths.mean() - mean_path_length)
        path_penalty = (path_lengths - path_mean).pow(2).mean()
        return path_penalty, path_mean.detach(), path_lengths

    def make_noise(self, batch, batch_size=None):
        if batch_size is None:
            batch_size = batch.size(0)
        if self.mixing_prob > 0 and random.random() < self.mixing_prob:
            return th.randn(2, batch_size, self.latent_size).type_as(batch).unbind(0)
        else:
            return [th.randn(batch_size, self.latent_size).type_as(batch)]

    def training_step(self, real_img, batch_idx, optimizer_idx):
        # real_img = real_img.half()
        log_dict = {}

        # train generator
        if optimizer_idx == 0:
            requires_grad(self.generator, True)
            requires_grad(self.discriminator, False)

            noise = self.make_noise(real_img)
            # print(real_img.dtype, noise[0].dtype, real_img.device)
            fake_img, _ = self.generator(noise)
            fake_pred = self.discriminator(fake_img)
            g_loss = self.g_nonsaturating_loss(fake_pred)

            log_dict["Generator"] = g_loss
            # log_dict["Spectral Norms/Generator"] = get_spectral_norms(self.generator)

            # print(g_loss)

            return OrderedDict({"loss": g_loss, "log": log_dict})

        # maybe regularize generator
        if optimizer_idx == 1:
            if batch_idx % self.g_reg_every == 0:
                path_batch_size = max(1, self.batch_size // self.path_batch_shrink)
                noise = self.make_noise(real_img, path_batch_size)
                fake_img, latents = self.generator(noise, return_latents=True)
                path_loss, self.mean_path_length, path_lengths = self.g_path_regularize(
                    fake_img, latents, self.mean_path_length.type_as(real_img)
                )
                weighted_path_loss = self.path_regularize * self.g_reg_every * path_loss
                if self.path_batch_shrink:
                    weighted_path_loss += 0 * fake_img[0, 0, 0, 0]

                log_dict["Path Length Regularization"] = path_loss
                log_dict["Mean Path Length"] = path_lengths.mean()

                return OrderedDict({"loss": weighted_path_loss, "log": log_dict})
            return OrderedDict({"loss": th.tensor(-69).type_as(real_img)})

        # train discriminator
        if optimizer_idx == 2:
            requires_grad(self.generator, False)
            requires_grad(self.discriminator, True)

            noise = self.make_noise(real_img)
            fake_img, _ = self.generator(noise)
            fake_pred = self.discriminator(fake_img)

            real_pred = self.discriminator(real_img)
            d_loss = self.d_logistic_loss(real_pred, fake_pred)

            log_dict["Discriminator"] = d_loss
            log_dict["Real Score"] = real_pred.mean()
            log_dict["Fake Score"] = fake_pred.mean()
            # log_dict["Spectral Norms/Discriminator"] = get_spectral_norms(self.discriminator)

            # print(d_loss)

            return OrderedDict({"loss": d_loss, "log": log_dict})

        # maybe regularize discriminator
        if optimizer_idx == 3:
            if batch_idx % self.d_reg_every == 0:
                real_img.requires_grad = True
                real_pred = self.discriminator(real_img)
                r1_loss = self.d_r1_loss(real_pred, real_img)
                weighted_r1_loss = self.r1 / 2 * r1_loss * self.d_reg_every + 0 * real_pred[0]

                log_dict["R1"] = r1_loss

                return OrderedDict({"loss": weighted_r1_loss, "log": log_dict})
            return OrderedDict({"loss": th.tensor(-69).type_as(real_img)})

    def backward(self, trainer, loss, optimizer, optimizer_idx):
        if optimizer_idx == 0:
            super(StyleGAN2, self).backward(trainer, loss, optimizer, optimizer_idx)
        if optimizer_idx == 1 and loss != -69:
            super(StyleGAN2, self).backward(trainer, loss, optimizer, optimizer_idx)
        if optimizer_idx == 2:
            super(StyleGAN2, self).backward(trainer, loss, optimizer, optimizer_idx)
        if optimizer_idx == 3 and loss != -69:
            super(StyleGAN2, self).backward(trainer, loss, optimizer, optimizer_idx)

    def optimizer_step(self, cur_epoch, batch_idx, optimizer, optimizer_idx, closure):
        if optimizer_idx == 0:
            super(StyleGAN2, self).optimizer_step(cur_epoch, batch_idx, optimizer, optimizer_idx, closure)
        if optimizer_idx == 1:
            if batch_idx % self.g_reg_every == 0:
                super(StyleGAN2, self).optimizer_step(cur_epoch, batch_idx, optimizer, optimizer_idx, closure)
            self.accumulate_g()
        if optimizer_idx == 2:
            super(StyleGAN2, self).optimizer_step(cur_epoch, batch_idx, optimizer, optimizer_idx, closure)
        if optimizer_idx == 3 and batch_idx % self.d_reg_every == 0:
            super(StyleGAN2, self).optimizer_step(cur_epoch, batch_idx, optimizer, optimizer_idx, closure)

    def prepare_data(self):
        validation.get_dataset_inception_features(self.train_dataloader(), self.name, self.size)

    def val_dataloader(self):
        return [[th.arange(0, 1)]]

    def validation_step(self, batch, batch_idx):
        # gc.collect()
        # th.cuda.empty_cache()
        # output = OrderedDict({"FID": th.tensor(-69).type_as(batch), "PPL": th.tensor(-69).type_as(batch)})
        # for task in batch:
        #     # if task == 1:
        #     output["FID"] = fid.validation_fid(
        #         self.g_ema.to(batch.device), self.val_batch_size, self.fid_n_sample, self.fid_truncation, self.name,
        #     )
        #     # if task == 0:
        #     output["PPL"] = ppl.validation_ppl(
        #         self.g_ema.to(batch.device),
        #         self.val_batch_size,
        #         self.ppl_n_sample,
        #         self.ppl_space,
        #         self.ppl_crop,
        #         self.latent_size,
        #     )
        return OrderedDict({"batch": batch})  # output

    def validation_epoch_end(self, outputs):
        batch = outputs[0]["batch"]
        gc.collect()
        th.cuda.empty_cache()
        val_fid = validation.fid(
            self.g_ema.to(batch.device), self.val_batch_size, self.fid_n_sample, self.fid_truncation, self.name,
        )["FID"]
        val_ppl = validation.ppl(
            self.g_ema.to(batch.device),
            self.val_batch_size,
            self.ppl_n_sample,
            self.ppl_space,
            self.ppl_crop,
            self.latent_size,
        )
        with th.no_grad():
            self.g_ema.eval()
            sample, _ = self.g_ema([self.sample_z.to(next(self.g_ema.parameters()).device)])
            grid = tv.utils.make_grid(
                sample, nrow=int(round(4.0 / 3 * self.n_sample ** 0.5)), normalize=True, range=(-1, 1)
            )
            self.logger.experiment.log(
                {"Generated Images EMA": [wandb.Image(grid, caption=f"Step {self.global_step}")]}
            )

            self.generator.eval()
            sample, _ = self.generator([self.sample_z.to(next(self.generator.parameters()).device)])
            grid = tv.utils.make_grid(
                sample, nrow=int(round(4.0 / 3 * self.n_sample ** 0.5)), normalize=True, range=(-1, 1)
            )
            self.logger.experiment.log({"Generated Images": [wandb.Image(grid, caption=f"Step {self.global_step}")]})
            self.generator.train()

        # val_fid = [score for score in outputs[0]["FID"] if score != -69][0]
        # val_ppl = [score for score in outputs[0]["PPL"] if score != -69][0]

        gc.collect()
        th.cuda.empty_cache()

        return {"val_loss": val_fid, "log": {"Validation/FID": val_fid, "Validation/PPL": val_ppl}}


if __name__ == "__main__":
    parser = argparse.ArgumentParser()

    # data options
    parser.add_argument("path", type=str)
    parser.add_argument("--vflip", type=bool, default=False)
    parser.add_argument("--hflip", type=bool, default=True)

    # training options
    parser.add_argument("--epochs", type=int, default=100)
    parser.add_argument("--batch_size", type=int, default=16)
    parser.add_argument("--checkpoint", type=str, default=None)

    # model options
    parser.add_argument("--latent_size", type=int, default=512)
    parser.add_argument("--n_mlp", type=int, default=8)
    parser.add_argument("--n_sample", type=int, default=32)
    parser.add_argument("--size", type=int, default=256)

    # loss options
    parser.add_argument("--r1", type=float, default=10)
    parser.add_argument("--path_regularize", type=float, default=2)
    parser.add_argument("--path_batch_shrink", type=int, default=2)
    parser.add_argument("--d_reg_every", type=int, default=16)
    parser.add_argument("--g_reg_every", type=int, default=4)
    parser.add_argument("--mixing_prob", type=float, default=0.9)
    parser.add_argument("--lr", type=float, default=0.002)
    parser.add_argument("--channel_multiplier", type=int, default=2)

    # validation / logging options
    parser.add_argument("--wandb", type=bool, default=True)
    parser.add_argument("--validation_interval", type=float, default=0.25)
    parser.add_argument("--val_batch_size", type=int, default=24)
    parser.add_argument("--fid_n_sample", type=int, default=10000)
    parser.add_argument("--fid_truncation", type=float, default=0.7)
    parser.add_argument("--ppl_space", choices=["z", "w"], default="w")
    parser.add_argument("--ppl_n_sample", type=int, default=5000)
    parser.add_argument("--ppl_crop", type=bool, default=False)

    # DevOps options
    parser.add_argument("--num_gpus", type=int, default=2)
    parser.add_argument("--cudnn_benchmark", type=bool, default=True)
    parser.add_argument("--distributed_backend", type=str, default="ddp")

    args = parser.parse_args()

    args.name = os.path.splitext(os.path.basename(args.path))[0]

    stylegan2 = StyleGAN2(args)
    stylegan2.prepare_data()
    stylegan2.train_dataloader()

    checkpoint_callback = pl.callbacks.ModelCheckpoint(
        filepath="/home/hans/modelzoo/maua-sg2/" + args.name + "-{epoch}-{val_loss:.0f}", save_top_k=10
    )

    wandb_logger = pl.loggers.WandbLogger(project="maua-stylegan")
    # print(wandb_logger.experiment)

    trainer = pl.Trainer(
        gpus=args.num_gpus,
        max_epochs=args.epochs,
        logger=wandb_logger,
        checkpoint_callback=checkpoint_callback,
        early_stop_callback=None,
        distributed_backend=args.distributed_backend,
        benchmark=args.cudnn_benchmark,
        val_check_interval=args.validation_interval,
        num_sanity_val_steps=0,
        terminate_on_nan=True,
        resume_from_checkpoint=args.checkpoint,
        amp_level="O2",
        precision=16,
    )
    trainer.fit(stylegan2)


================================================
FILE: lookahead_minimax.py
================================================
from collections import defaultdict

import torch
from torch.optim.optimizer import Optimizer


class LookaheadMinimax(Optimizer):
    r"""
    A PyTorch implementation of the lookahead wrapper for GANs.

    This optimizer performs the lookahead step on both the discriminator and generator optimizers after the generator's
    optimizer takes a step. This ensures that joint minimax lookahead is used rather than alternating minimax lookahead
    (which would result from simply applying the original Lookahead Optimizer to both networks separately).

    Lookahead Minimax Optimizer: https://arxiv.org/abs/2006.14567
    Lookahead Optimizer: https://arxiv.org/abs/1907.08610
    """

    def __init__(self, G_optimizer, D_optimizer, la_steps=5, la_alpha=0.5, pullback_momentum="none", accumulate=1):
        """
        G_optimizer: generator optimizer
        D_optimizer: discriminator optimizer
        la_steps (int): number of lookahead steps
        la_alpha (float): linear interpolation factor. 1.0 recovers the inner optimizer.
        pullback_momentum (str): change to inner optimizer momentum on interpolation update
        acumulate (int): number of gradient accumulation steps
        """
        self.G_optimizer = G_optimizer
        self.D_optimizer = D_optimizer

        self._la_step = 0  # counter for inner optimizer
        self.la_alpha = la_alpha
        self._total_la_steps = la_steps * accumulate
        self._la_steps = la_steps

        pullback_momentum = pullback_momentum.lower()
        assert pullback_momentum in ["reset", "pullback", "none"]
        self.pullback_momentum = pullback_momentum

        self.state = defaultdict(dict)

        # Cache the current optimizer parameters
        for group in G_optimizer.param_groups:
            for p in group["params"]:
                param_state = self.state[p]
                param_state["cached_G_params"] = torch.zeros_like(p.data)
                param_state["cached_G_params"].copy_(p.data)
                if self.pullback_momentum == "pullback":
                    param_state["cached_G_mom"] = torch.zeros_like(p.data)

        for group in D_optimizer.param_groups:
            for p in group["params"]:
                param_state = self.state[p]
                param_state["cached_D_params"] = torch.zeros_like(p.data)
                param_state["cached_D_params"].copy_(p.data)
                if self.pullback_momentum == "pullback":
                    param_state["cached_D_mom"] = torch.zeros_like(p.data)

    def __getstate__(self):
        return {
            "state": self.state,
            "G_optimizer": self.G_optimizer,
            "D_optimizer": self.D_optimizer,
            "la_alpha": self.la_alpha,
            "_la_step": self._la_step,
            "_total_la_steps": self._la_steps,
            "pullback_momentum": self.pullback_momentum,
        }

    def zero_grad(self):
        self.G_optimizer.zero_grad()

    def get_la_step(self):
        return self._la_step

    def state_dict(self):
        return self.G_optimizer.state_dict()

    def load_state_dict(self, G_state_dict, D_state_dict):
        self.G_optimizer.load_state_dict(G_state_dict)
        self.D_optimizer.load_state_dict(D_state_dict)

        # Cache the current optimizer parameters
        for group in self.G_optimizer.param_groups:
            for p in group["params"]:
                param_state = self.state[p]
                param_state["cached_G_params"] = torch.zeros_like(p.data)
                param_state["cached_G_params"].copy_(p.data)
                if self.pullback_momentum == "pullback":
                    param_state["cached_G_mom"] = self.G_optimizer.state[p]["momentum_buffer"]

        for group in self.D_optimizer.param_groups:
            for p in group["params"]:
                param_state = self.state[p]
                param_state["cached_D_params"] = torch.zeros_like(p.data)
                param_state["cached_D_params"].copy_(p.data)
                if self.pullback_momentum == "pullback":
                    param_state["cached_D_mom"] = self.D_optimizer.state[p]["momentum_buffer"]

    def _backup_and_load_cache(self):
        """
        Useful for performing evaluation on the slow weights (which typically generalize better)
        """
        for group in self.G_optimizer.param_groups:
            for p in group["params"]:
                param_state = self.state[p]
                param_state["backup_G_params"] = torch.zeros_like(p.data)
                param_state["backup_G_params"].copy_(p.data)
                p.data.copy_(param_state["cached_G_params"])

        for group in self.D_optimizer.param_groups:
            for p in group["params"]:
                param_state = self.state[p]
                param_state["backup_D_params"] = torch.zeros_like(p.data)
                param_state["backup_D_params"].copy_(p.data)
                p.data.copy_(param_state["cached_D_params"])

    def _clear_and_load_backup(self):
        for group in self.G_optimizer.param_groups:
            for p in group["params"]:
                param_state = self.state[p]
                p.data.copy_(param_state["backup_G_params"])
                del param_state["backup_G_params"]

        for group in self.D_optimizer.param_groups:
            for p in group["params"]:
                param_state = self.state[p]
                p.data.copy_(param_state["backup_D_params"])
                del param_state["backup_D_params"]

    @property
    def param_groups(self):
        return self.G_optimizer.param_groups

    def step(self, closure=None):
        """
        Performs a single Lookahead optimization step on BOTH optimizers after the generator's optimizer step.

        This allows the discriminator's optimizer to take more steps when using a higher step ratio and still have the
        lookahead step being performed once after k generator steps. This also ensures the optimizers are updated with
        the lookahead step simultaneously, rather than in alternating fashion.

        Arguments:
            closure (callable, optional): A closure that reevaluates the model
                and returns the loss.
        """
        loss = self.G_optimizer.step(closure)
        self._la_step += 1

        if self._la_step >= self._total_la_steps:
            with torch.cuda.amp.autocast(enabled=False):
                self._la_step = 0

                # Lookahead and cache the current generator optimizer parameters
                for group in self.G_optimizer.param_groups:
                    for p in group["params"]:
                        param_state = self.state[p]
                        p.data.mul_(self.la_alpha).add_(1.0 - self.la_alpha, param_state["cached_G_params"])
                        param_state["cached_G_params"].copy_(p.data)

                        if self.pullback_momentum == "pullback":
                            internal_momentum = self.G_optimizer.state[p]["momentum_buffer"]
                            self.G_optimizer.state[p]["momentum_buffer"] = internal_momentum.mul_(self.la_alpha).add_(
                                1.0 - self.la_alpha, param_state["cached_G_mom"]
                            )
                            param_state["cached_G_mom"] = self.G_optimizer.state[p]["momentum_buffer"]
                        elif self.pullback_momentum == "reset":
                            self.G_optimizer.state[p]["momentum_buffer"] = torch.zeros_like(p.data)

                # Lookahead and cache the current discriminator optimizer parameters
                for group in self.D_optimizer.param_groups:
                    for p in group["params"]:
                        param_state = self.state[p]
                        p.data.mul_(self.la_alpha).add_(1.0 - self.la_alpha, param_state["cached_D_params"])
                        param_state["cached_D_params"].copy_(p.data)

                        if self.pullback_momentum == "pullback":
                            internal_momentum = self.D_optimizer.state[p]["momentum_buffer"]
                            self.D_optimizer.state[p]["momentum_buffer"] = internal_momentum.mul_(self.la_alpha).add_(
                                1.0 - self.la_alpha, param_state["cached_D_mom"]
                            )
                            param_state["cached_D_mom"] = self.optimizer.state[p]["momentum_buffer"]
                        elif self.pullback_momentum == "reset":
                            self.D_optimizer.state[p]["momentum_buffer"] = torch.zeros_like(p.data)

        return loss


================================================
FILE: lucidrains.py
================================================
import json, time, pickle, argparse
from math import floor, log2, sqrt
from random import random
from shutil import rmtree
from functools import partial
from datetime import datetime
import multiprocessing
from PIL import Image
from pathlib import Path
from retry.api import retry_call
from tqdm import tqdm

import numpy as np
from scipy import linalg

import torch
from torch import nn
from torch.utils import data
import torch.nn.functional as F

from torch_optimizer import DiffGrad
from torch.autograd import grad as torch_grad

import torchvision
from torchvision import transforms

from vector_quantize_pytorch import VectorQuantize
from linear_attention_transformer import ImageLinearAttention
from contrastive_learner import ContrastiveLearner, RandomApply

from kornia import augmentation as augs
from kornia import filters

import validation
from validation.inception import InceptionV3
from validation import lpips

import wandb

try:
    from apex import amp

    APEX_AVAILABLE = True
except:
    APEX_AVAILABLE = False

assert torch.cuda.is_available(), "You need to have an Nvidia GPU with CUDA installed."

num_cores = multiprocessing.cpu_count()

# constants

EXTS = ["jpg", "png"]
EPS = 1e-8

# helper classes


class NanException(Exception):
    pass


class EMA:
    def __init__(self, beta):
        super().__init__()
        self.beta = beta

    def update_average(self, old, new):
        if old is None:
            return new
        return old * self.beta + (1 - self.beta) * new


class Flatten(nn.Module):
    def forward(self, x):
        return x.reshape(x.shape[0], -1)


class Residual(nn.Module):
    def __init__(self, fn):
        super().__init__()
        self.fn = fn

    def forward(self, x):
        return self.fn(x) + x


class Rezero(nn.Module):
    def __init__(self, fn):
        super().__init__()
        self.fn = fn
        self.g = nn.Parameter(torch.zeros(1))

    def forward(self, x):
        return self.fn(x) * self.g


class PermuteToFrom(nn.Module):
    def __init__(self, fn):
        super().__init__()
        self.fn = fn

    def forward(self, x):
        x = x.permute(0, 2, 3, 1)
        out, loss = self.fn(x)
        out = out.permute(0, 3, 1, 2)
        return out, loss


# helpers


def default(value, d):
    return d if value is None else value


def cycle(iterable):
    while True:
        for i in iterable:
            yield i


def cast_list(el):
    return el if isinstance(el, list) else [el]


def is_empty(t):
    if isinstance(t, torch.Tensor):
        return t.nelement() == 0
    return t is None


def raise_if_nan(t):
    if torch.isnan(t):
        raise NanException


def loss_backwards(fp16, loss, optimizer, **kwargs):
    if fp16:
        with amp.scale_loss(loss, optimizer) as scaled_loss:
            scaled_loss.backward(**kwargs)
    else:
        loss.backward(**kwargs)


def gradient_penalty(images, output, weight=10):
    batch_size = images.shape[0]
    gradients = torch_grad(
        outputs=output,
        inputs=images,
        grad_outputs=torch.ones(output.size()).cuda(),
        create_graph=True,
        retain_graph=True,
        only_inputs=True,
    )[0]

    gradients = gradients.view(batch_size, -1)
    return weight * ((gradients.norm(2, dim=1) - 1) ** 2).mean()


def noise(n, latent_dim):
    return torch.randn(n, latent_dim).cuda()


def noise_list(n, layers, latent_dim):
    return [(noise(n, latent_dim), layers)]


def mixed_list(n, layers, latent_dim):
    tt = int(torch.rand(()).numpy() * layers)
    return noise_list(n, tt, latent_dim) + noise_list(n, layers - tt, latent_dim)


def latent_to_w(style_vectorizer, latent_descr):
    return [(style_vectorizer(z), num_layers) for z, num_layers in latent_descr]


def image_noise(n, im_size):
    return torch.FloatTensor(n, im_size, im_size, 1).uniform_(0.0, 1.0).cuda()


def leaky_relu(p=0.2):
    return nn.LeakyReLU(p, inplace=True)


def evaluate_in_chunks(max_batch_size, model, *args):
    split_args = list(zip(*list(map(lambda x: x.split(max_batch_size, dim=0), args))))
    chunked_outputs = [model(*i) for i in split_args]
    if len(chunked_outputs) == 1:
        return chunked_outputs[0]
    return torch.cat(chunked_outputs, dim=0)


def styles_def_to_tensor(styles_def):
    return torch.cat([t[:, None, :].expand(-1, n, -1) for t, n in styles_def], dim=1)


def set_requires_grad(model, bool):
    for p in model.parameters():
        p.requires_grad = bool


# dataset


def convert_rgb_to_transparent(image):
    if image.mode == "RGB":
        return image.convert("RGBA")
    return image


def convert_transparent_to_rgb(image):
    if image.mode == "RGBA":
        return image.convert("RGB")
    return image


class expand_greyscale(object):
    def __init__(self, num_channels):
        self.num_channels = num_channels

    def __call__(self, tensor):
        return tensor.expand(self.num_channels, -1, -1)


def resize_to_minimum_size(min_size, image):
    if max(*image.size) < min_size:
        return torchvision.transforms.functional.resize(image, min_size)
    return image


class Dataset(data.Dataset):
    def __init__(self, folder, image_size, transparent=False):
        super().__init__()
        self.folder = folder
        self.image_size = image_size
        self.paths = [p for ext in EXTS for p in Path(f"{folder}").glob(f"**/*.{ext}")]

        convert_image_fn = convert_transparent_to_rgb if not transparent else convert_rgb_to_transparent
        num_channels = 3 if not transparent else 4

        self.transform = transforms.Compose(
            [
                transforms.Lambda(convert_image_fn),
                transforms.Lambda(partial(resize_to_minimum_size, image_size)),
                transforms.RandomHorizontalFlip(),
                transforms.Resize(image_size),
                transforms.CenterCrop(image_size),
                transforms.ToTensor(),
                transforms.Lambda(expand_greyscale(num_channels)),
            ]
        )

    def __len__(self):
        return len(self.paths)

    def __getitem__(self, index):
        path = self.paths[index]
        img = Image.open(path)
        return self.transform(img)


# stylegan2 classes


class StyleVectorizer(nn.Module):
    def __init__(self, emb, depth):
        super().__init__()

        layers = []
        for i in range(depth):
            layers.extend([nn.Linear(emb, emb), leaky_relu()])

        self.net = nn.Sequential(*layers)

    def forward(self, x):
        return self.net(x)


class RGBBlock(nn.Module):
    def __init__(self, latent_dim, input_channel, upsample, rgba=False):
        super().__init__()
        self.input_channel = input_channel
        self.to_style = nn.Linear(latent_dim, input_channel)

        out_filters = 3 if not rgba else 4
        self.conv = Conv2DMod(input_channel, out_filters, 1, demod=False)

        self.upsample = nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False) if upsample else None

    def forward(self, x, prev_rgb, istyle):
        b, c, h, w = x.shape
        style = self.to_style(istyle)
        x = self.conv(x, style)

        if prev_rgb is not None:
            x = x + prev_rgb

        if self.upsample is not None:
            x = self.upsample(x)

        return x


class Conv2DMod(nn.Module):
    def __init__(self, in_chan, out_chan, kernel, demod=True, stride=1, dilation=1, **kwargs):
        super().__init__()
        self.filters = out_chan
        self.demod = demod
        self.kernel = kernel
        self.stride = stride
        self.dilation = dilation
        self.weight = nn.Parameter(torch.randn((out_chan, in_chan, kernel, kernel)))
        nn.init.kaiming_normal_(self.weight, a=0, mode="fan_in", nonlinearity="leaky_relu")

    def _get_same_padding(self, size, kernel, dilation, stride):
        return ((size - 1) * (stride - 1) + dilation * (kernel - 1)) // 2

    def forward(self, x, y):
        b, c, h, w = x.shape

        w1 = y[:, None, :, None, None]
        w2 = self.weight[None, :, :, :, :]
        weights = w2 * (w1 + 1)

        if self.demod:
            d = torch.rsqrt((weights ** 2).sum(dim=(2, 3, 4), keepdim=True) + EPS)
            weights = weights * d

        x = x.reshape(1, -1, h, w)

        _, _, *ws = weights.shape
        weights = weights.reshape(b * self.filters, *ws)

        padding = self._get_same_padding(h, self.kernel, self.dilation, self.stride)
        x = F.conv2d(x, weights, padding=padding, groups=b)

        x = x.reshape(-1, self.filters, h, w)
        return x


class GeneratorBlock(nn.Module):
    def __init__(self, latent_dim, input_channels, filters, upsample=True, upsample_rgb=True, rgba=False):
        super().__init__()
        self.upsample = nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False) if upsample else None

        self.to_style1 = nn.Linear(latent_dim, input_channels)
        self.to_noise1 = nn.Linear(1, filters)
        self.conv1 = Conv2DMod(input_channels, filters, 3)

        self.to_style2 = nn.Linear(latent_dim, filters)
        self.to_noise2 = nn.Linear(1, filters)
        self.conv2 = Conv2DMod(filters, filters, 3)

        self.activation = leaky_relu()
        self.to_rgb = RGBBlock(latent_dim, filters, upsample_rgb, rgba)

    def forward(self, x, prev_rgb, istyle, inoise):
        if self.upsample is not None:
            x = self.upsample(x)

        inoise = inoise[:, : x.shape[2], : x.shape[3], :]
        noise1 = self.to_noise1(inoise).permute((0, 3, 2, 1))
        noise2 = self.to_noise2(inoise).permute((0, 3, 2, 1))

        style1 = self.to_style1(istyle)
        x = self.conv1(x, style1)
        x = self.activation(x + noise1)

        style2 = self.to_style2(istyle)
        x = self.conv2(x, style2)
        x = self.activation(x + noise2)

        rgb = self.to_rgb(x, prev_rgb, istyle)
        return x, rgb


class DiscriminatorBlock(nn.Module):
    def __init__(self, input_channels, filters, downsample=True):
        super().__init__()
        self.conv_res = nn.Conv2d(input_channels, filters, 1)

        self.net = nn.Sequential(
            nn.Conv2d(input_channels, filters, 3, padding=1),
            leaky_relu(),
            nn.Conv2d(filters, filters, 3, padding=1),
            leaky_relu(),
        )

        self.downsample = nn.Conv2d(filters, filters, 3, padding=1, stride=2) if downsample else None

    def forward(self, x):
        res = self.conv_res(x)
        x = self.net(x)
        x = x + res
        if self.downsample is not None:
            x = self.downsample(x)
        return x


class Generator(nn.Module):
    def __init__(self, image_size, latent_dim, network_capacity=16, transparent=False, attn_layers=[]):
        super().__init__()
        self.image_size = image_size
        self.latent_dim = latent_dim
        self.num_layers = int(log2(image_size) - 1)

        init_channels = 4 * network_capacity
        self.initial_block = nn.Parameter(torch.randn((init_channels, 4, 4)))
        filters = [init_channels] + [network_capacity * (2 ** (i + 1)) for i in range(self.num_layers)][::-1]
        in_out_pairs = zip(filters[0:-1], filters[1:])

        self.blocks = nn.ModuleList([])
        self.attns = nn.ModuleList([])

        for ind, (in_chan, out_chan) in enumerate(in_out_pairs):
            not_first = ind != 0
            not_last = ind != (self.num_layers - 1)
            num_layer = self.num_layers - ind

            attn_fn = (
                nn.Sequential(*[Residual(Rezero(ImageLinearAttention(in_chan))) for _ in range(2)])
                if num_layer in attn_layers
                else None
            )

            self.attns.append(attn_fn)

            block = GeneratorBlock(
                latent_dim, in_chan, out_chan, upsample=not_first, upsample_rgb=not_last, rgba=transparent
            )
            self.blocks.append(block)

    def forward(self, styles, input_noise):
        batch_size = styles.shape[0]
        image_size = self.image_size
        x = self.initial_block.expand(batch_size, -1, -1, -1)
        styles = styles.transpose(0, 1)

        rgb = None
        for style, block, attn in zip(styles, self.blocks, self.attns):
            if attn is not None:
                x = attn(x)
            x, rgb = block(x, rgb, style, input_noise)

        return rgb


class Discriminator(nn.Module):
    def __init__(
        self, image_size, network_capacity=16, fq_layers=[], fq_dict_size=256, attn_layers=[], transparent=False
    ):
        super().__init__()
        num_layers = int(log2(image_size) - 1)
        num_init_filters = 3 if not transparent else 4

        blocks = []
        filters = [num_init_filters] + [(network_capacity) * (2 ** i) for i in range(num_layers + 1)]
        chan_in_out = list(zip(filters[0:-1], filters[1:]))

        blocks = []
        quantize_blocks = []
        attn_blocks = []

        for ind, (in_chan, out_chan) in enumerate(chan_in_out):
            num_layer = ind + 1
            is_not_last = ind != (len(chan_in_out) - 1)

            block = DiscriminatorBlock(in_chan, out_chan, downsample=is_not_last)
            blocks.append(block)

            attn_fn = (
                nn.Sequential(*[Residual(Rezero(ImageLinearAttention(out_chan))) for _ in range(2)])
                if num_layer in attn_layers
                else None
            )

            attn_blocks.append(attn_fn)

            quantize_fn = PermuteToFrom(VectorQuantize(out_chan, fq_dict_size)) if num_layer in fq_layers else None
            quantize_blocks.append(quantize_fn)

        self.blocks = nn.ModuleList(blocks)
        self.attn_blocks = nn.ModuleList(attn_blocks)
        self.quantize_blocks = nn.ModuleList(quantize_blocks)

        latent_dim = 2 * 2 * filters[-1]

        self.flatten = Flatten()
        self.to_logit = nn.Linear(latent_dim, 1)

    def forward(self, x):
        b, *_ = x.shape

        quantize_loss = torch.zeros(1).to(x)

        for (block, attn_block, q_block) in zip(self.blocks, self.attn_blocks, self.quantize_blocks):
            x = block(x)

            if attn_block is not None:
                x = attn_block(x)

            if q_block is not None:
                x, loss = q_block(x)
                quantize_loss += loss

        x = self.flatten(x)
        x = self.to_logit(x)
        return x.squeeze(), quantize_loss


class StyleGAN2(nn.Module):
    def __init__(
        self,
        image_size,
        latent_dim=512,
        style_depth=8,
        network_capacity=16,
        transparent=False,
        fp16=False,
        cl_reg=False,
        augment_fn=None,
        steps=1,
        lr=1e-4,
        fq_layers=[],
        fq_dict_size=256,
        attn_layers=[],
    ):
        super().__init__()
        self.lr = lr
        self.steps = steps
        self.ema_updater = EMA(0.995)

        self.S = StyleVectorizer(latent_dim, style_depth)
        self.G = Generator(image_size, latent_dim, network_capacity, transparent=transparent, attn_layers=attn_layers)
        self.D = Discriminator(
            image_size,
            network_capacity,
            fq_layers=fq_layers,
            fq_dict_size=fq_dict_size,
            attn_layers=attn_layers,
            transparent=transparent,
        )

        self.SE = StyleVectorizer(latent_dim, style_depth)
        self.GE = Generator(image_size, latent_dim, network_capacity, transparent=transparent, attn_layers=attn_layers)

        set_requires_grad(self.SE, False)
        set_requires_grad(self.GE, False)

        generator_params = list(self.G.parameters()) + list(self.S.parameters())
        self.G_opt = DiffGrad(generator_params, lr=self.lr, betas=(0.5, 0.9))
        self.D_opt = DiffGrad(self.D.parameters(), lr=self.lr, betas=(0.5, 0.9))

        self._init_weights()
        self.reset_parameter_averaging()

        self.cuda()

        if fp16:
            (self.S, self.G, self.D, self.SE, self.GE), (self.G_opt, self.D_opt) = amp.initialize(
                [self.S, self.G, self.D, self.SE, self.GE], [self.G_opt, self.D_opt], opt_level="O2"
            )

        # experimental contrastive loss discriminator regularization
        if augment_fn is not None:
            self.augment_fn = augment_fn
        else:
            self.augment_fn = nn.Sequential(
                nn.ReflectionPad2d(int((sqrt(2) - 1) * image_size / 4)),
                RandomApply(augs.ColorJitter(0.8, 0.8, 0.8, 0.2), p=0.7),
                augs.RandomGrayscale(p=0.2),
                augs.RandomHorizontalFlip(),
                RandomApply(augs.RandomAffine(degrees=0, translate=(0.25, 0.25), shear=(15, 15)), p=0.3),
                RandomApply(
                    nn.Sequential(augs.RandomRotation(180), augs.CenterCrop(size=(image_size, image_size))), p=0.2
                ),
                augs.RandomResizedCrop(size=(image_size, image_size)),
                RandomApply(filters.GaussianBlur2d((3, 3), (1.5, 1.5)), p=0.1),
                RandomApply(augs.RandomErasing(), p=0.1),
            )

        self.D_cl = (
            ContrastiveLearner(self.D, image_size, augment_fn=self.augment_fn, fp16=fp16, hidden_layer="flatten")
            if cl_reg
            else None
        )

        # self.S, self.G, self.D, self.SE, self.GE = (
        #     nn.DataParallel(self.S),
        #     nn.DataParallel(self.G),
        #     nn.DataParallel(self.D),
        #     nn.DataParallel(self.SE),
        #     nn.DataParallel(self.GE),
        # )

    def _init_weights(self):
        for m in self.modules():
            if type(m) in {nn.Conv2d, nn.Linear}:
                nn.init.kaiming_normal_(m.weight, a=0, mode="fan_in", nonlinearity="leaky_relu")

        for block in self.G.blocks:
            nn.init.zeros_(block.to_noise1.weight)
            nn.init.zeros_(block.to_noise2.weight)
            nn.init.zeros_(block.to_noise1.bias)
            nn.init.zeros_(block.to_noise2.bias)

    def EMA(self):
        def update_moving_average(ma_model, current_model):
            for current_params, ma_params in zip(current_model.parameters(), ma_model.parameters()):
                old_weight, up_weight = ma_params.data, current_params.data
                ma_params.data = self.ema_updater.update_average(old_weight, up_weight)

        update_moving_average(self.SE, self.S)
        update_moving_average(self.GE, self.G)

    def reset_parameter_averaging(self):
        self.SE.load_state_dict(self.S.state_dict())
        self.GE.load_state_dict(self.G.state_dict())

    def forward(self, x):
        return x


class Trainer:
    def __init__(
        self,
        name,
        results_dir,
        models_dir,
        image_size,
        network_capacity,
        transparent=False,
        batch_size=4,
        mixed_prob=0.9,
        gradient_accumulate_every=1,
        lr=2e-4,
        num_workers=None,
        save_every=1000,
        trunc_psi=0.6,
        fp16=False,
        cl_reg=False,
        fq_layers=[],
        fq_dict_size=256,
        attn_layers=[],
        fid_n_sample=5000,
        ppl_n_sample=2500,
        *args,
        **kwargs,
    ):
        self.GAN_params = [args, kwargs]
        self.GAN = None

        self.name = name
        self.results_dir = Path(results_dir)
        self.models_dir = Path(models_dir)
        self.config_path = self.models_dir / name / ".config.json"

        assert log2(image_size).is_integer(), "image size must be a power of 2 (64, 128, 256, 512, 1024)"
        self.image_size = image_size
        self.network_capacity = network_capacity
        self.transparent = transparent
        self.fq_layers = cast_list(fq_layers)
        self.fq_dict_size = fq_dict_size

        self.attn_layers = cast_list(attn_layers)

        self.lr = lr
        self.batch_size = batch_size
        self.num_workers = num_workers
        self.mixed_prob = mixed_prob

        self.save_every = save_every
        self.steps = 0

        self.av = None
        self.trunc_psi = trunc_psi

        self.pl_mean = 0

        self.fid_n_sample = fid_n_sample
        self.ppl_n_sample = ppl_n_sample

        self.gradient_accumulate_every = gradient_accumulate_every

        assert not fp16 or fp16 and APEX_AVAILABLE, "Apex is not available for you to use mixed precision training"
        self.fp16 = fp16

        self.cl_reg = cl_reg

        self.d_loss = 0
        self.g_loss = 0
        self.last_gp_loss = 0
        self.last_cr_loss = 0
        self.q_loss = 0

        self.pl_length_ma = EMA(0.99)
        self.init_folders()

        self.loader = None

    def init_GAN(self):
        args, kwargs = self.GAN_params
        self.GAN = StyleGAN2(
            lr=self.lr,
            image_size=self.image_size,
            network_capacity=self.network_capacity,
            transparent=self.transparent,
            fq_layers=self.fq_layers,
            fq_dict_size=self.fq_dict_size,
            attn_layers=self.attn_layers,
            fp16=self.fp16,
            cl_reg=self.cl_reg,
            *args,
            **kwargs,
        )

    def write_config(self):
        self.config_path.write_text(json.dumps(self.config()))

    def load_config(self):
        config = self.config() if not self.config_path.exists() else json.loads(self.config_path.read_text())
        self.image_size = config["image_size"]
        self.network_capacity = config["network_capacity"]
        self.transparent = config["transparent"]
        self.fq_layers = config["fq_layers"]
        self.fq_dict_size = config["fq_dict_size"]
        self.attn_layers = config.pop("attn_layers", [])
        del self.GAN
        self.init_GAN()

    def config(self):
        return {
            "image_size": self.image_size,
            "network_capacity": self.network_capacity,
            "transparent": self.transparent,
            "fq_layers": self.fq_layers,
            "fq_dict_size": self.fq_dict_size,
            "attn_layers": self.attn_layers,
        }

    def set_data_src(self, folder):
        self.dataset = Dataset(folder, self.image_size, transparent=self.transparent)
        self.loader = cycle(
            data.DataLoader(
                self.dataset,
                num_workers=default(self.num_workers, num_cores),
                batch_size=self.batch_size,
                drop_last=True,
                shuffle=True,
                pin_memory=True,
            )
        )
        validation.get_dataset_inception_features(self.loader, self.name, self.image_size)

    def train(self):
        assert (
            self.loader is not None
        ), "You must first initialize the data source with `.set_data_src(<folder of images>)`"

        if self.GAN is None:
            self.init_GAN()

        self.GAN.train()
        total_disc_loss = torch.tensor(0.0).cuda()
        total_gen_loss = torch.tensor(0.0).cuda()

        batch_size = self.batch_size

        image_size = self.GAN.G.image_size
        latent_dim = self.GAN.G.latent_dim
        num_layers = self.GAN.G.num_layers

        apply_gradient_penalty = self.steps % 4 == 0
        apply_path_penalty = self.steps % 32 == 0
        apply_cl_reg_to_generated = self.steps > 20000

        log_dict = {"Divergence": 0, "Quantize": 0, "Generator": 0}
        if apply_gradient_penalty:
            log_dict["R1"] = 0
        if apply_path_penalty:
            log_dict["Path Length"] = 0

        backwards = partial(loss_backwards, self.fp16)

        if self.GAN.D_cl is not None:
            self.GAN.D_opt.zero_grad()

            if apply_cl_reg_to_generated:
                for i in range(self.gradient_accumulate_every):
                    get_latents_fn = mixed_list if random() < self.mixed_prob else noise_list
                    style = get_latents_fn(batch_size, num_layers, latent_dim)
                    noise = image_noise(batch_size, image_size)

                    w_space = latent_to_w(self.GAN.S, style)
                    w_styles = styles_def_to_tensor(w_space)

                    generated_images = self.GAN.G(w_styles, noise)
                    self.GAN.D_cl(generated_images.clone().detach(), accumulate=True)

            for i in range(self.gradient_accumulate_every):
                image_batch = next(self.loader).cuda()
                self.GAN.D_cl(image_batch, accumulate=True)

            loss = self.GAN.D_cl.calculate_loss()
            self.last_cr_loss = loss.clone().detach().item()
            log_dict["Consistency"] = self.last_cr_loss
            backwards(loss, self.GAN.D_opt)

            self.GAN.D_opt.step()

        # train discriminator

        avg_pl_length = self.pl_mean
        self.GAN.D_opt.zero_grad()

        for i in range(self.gradient_accumulate_every):
            get_latents_fn = mixed_list if random() < self.mixed_prob else noise_list
            style = get_latents_fn(batch_size, num_layers, latent_dim)
            noise = image_noise(batch_size, image_size)

            w_space = latent_to_w(self.GAN.S, style)
            w_styles = styles_def_to_tensor(w_space)

            generated_images = self.GAN.G(w_styles, noise)
            fake_output, fake_q_loss = self.GAN.D(generated_images.clone().detach())

            image_batch = next(self.loader).cuda()
            image_batch.requires_grad_()
            real_output, real_q_loss = self.GAN.D(image_batch)

            divergence = (F.relu(1 + real_output) + F.relu(1 - fake_output)).mean()
            disc_loss = divergence
            log_dict["Divergence"] += divergence / self.gradient_accumulate_every

            quantize_loss = (fake_q_loss + real_q_loss).mean()
            self.q_loss = float(quantize_loss.detach().item())
            log_dict["Quantize"] += self.q_loss / self.gradient_accumulate_every

            disc_loss = disc_loss + quantize_loss

            if apply_gradient_penalty:
                gp = gradient_penalty(image_batch, real_output)
                self.last_gp_loss = gp.clone().detach().item()
                disc_loss = disc_loss + gp
                log_dict["R1"] += gp / self.gradient_accumulate_every

            disc_loss = disc_loss / self.gradient_accumulate_every
            disc_loss.register_hook(raise_if_nan)
            backwards(disc_loss, self.GAN.D_opt)

            total_disc_loss += divergence.detach().item() / self.gradient_accumulate_every

        self.d_loss = float(total_disc_loss)
        log_dict["Discriminator"] = self.d_loss
        self.GAN.D_opt.step()

        # train generator

        self.GAN.G_opt.zero_grad()
        for i in range(self.gradient_accumulate_every):
            style = get_latents_fn(batch_size, num_layers, latent_dim)
            noise = image_noise(batch_size, image_size)

            w_space = latent_to_w(self.GAN.S, style)
            w_styles = styles_def_to_tensor(w_space)

            generated_images = self.GAN.G(w_styles, noise)
            fake_output, _ = self.GAN.D(generated_images)
            loss = fake_output.mean()
            gen_loss = loss
            log_dict["Generator"] += gen_loss / self.gradient_accumulate_every

            if apply_path_penalty:
                std = 0.1 / (w_styles.std(dim=0, keepdim=True) + EPS)
                w_styles_2 = w_styles + torch.randn(w_styles.shape).cuda() / (std + EPS)
                pl_images = self.GAN.G(w_styles_2, noise)
                pl_lengths = ((pl_images - generated_images) ** 2).mean(dim=(1, 2, 3))
                avg_pl_length = np.mean(pl_lengths.detach().cpu().numpy())

                if not is_empty(self.pl_mean):
                    pl_loss = ((pl_lengths - self.pl_mean) ** 2).mean()
                    log_dict["Path Length"] += pl_loss / self.gradient_accumulate_every
                    if not torch.isnan(pl_loss):
                        gen_loss = gen_loss + pl_loss

            gen_loss = gen_loss / self.gradient_accumulate_every
            gen_loss.register_hook(raise_if_nan)
            backwards(gen_loss, self.GAN.G_opt)

            total_gen_loss += loss.detach().item() / self.gradient_accumulate_every

        self.g_loss = float(total_gen_loss)
        self.GAN.G_opt.step()

        # calculate moving averages
        if apply_path_penalty and not np.isnan(avg_pl_length):
            self.pl_mean = self.pl_length_ma.update_average(self.pl_mean, avg_pl_length)
            log_dict["Mean Path Length"] = self.pl_mean

        if self.steps % 10 == 0 and self.steps > 20000:
            self.GAN.EMA()

        if self.steps <= 25000 and self.steps % 1000 == 2:
            self.GAN.reset_parameter_averaging()

        # save from NaN errors
        checkpoint_num = floor(self.steps / self.save_every)

        if any(torch.isnan(l) for l in (total_gen_loss, total_disc_loss)):
            print(f"NaN detected for generator or discriminator. Loading from checkpoint #{checkpoint_num}")
            self.load(checkpoint_num)
            raise NanException

        # periodically save results
        if self.steps % self.save_every == 0:
            self.save(checkpoint_num)

        if self.steps % 1000 == 0 or (self.steps % 100 == 0 and self.steps < 2500):
            self.evaluate(floor(self.steps / 1000))

        if self.steps % 1000 == 0:
            start_time = time.time()

            PBAR.set_description((f"Calculating FID..."))
            fid, density, coverage = self.calculate_fid()
            log_dict["Evaluation/FID"] = fid
            log_dict["Evaluation/Density"] = density
            log_dict["Evaluation/Coverage"] = coverage

            PBAR.set_description((f"Calculating PPL..."))
            ppl = self.calculate_ppl()

            PBAR.set_description(
                (
                    f"FID: {fid:.4f}; Density: {density:.4f}; Coverage: {coverage:.4f}; PPL: {ppl:.4f} in {time.time() - start_time:.1f}s"
                )
            )
            log_dict["Evaluation/PPL"] = ppl

        wandb.log(log_dict)

        self.steps += 1
        self.av = None

    @torch.no_grad()
    def evaluate(self, num=0, num_image_tiles=8):
        self.GAN.eval()
        ext = "jpg" if not self.transparent else "png"
        num_rows = num_image_tiles

        latent_dim = self.GAN.G.latent_dim
        image_size = self.GAN.G.image_size
        num_layers = self.GAN.G.num_layers

        # latents and noise
        latents = noise_list(num_rows ** 2, num_layers, latent_dim)
        n = image_noise(num_rows ** 2, image_size)

        # regular
        generated_images = self.generate_truncated(self.GAN.S, self.GAN.G, latents, n, trunc_psi=self.trunc_psi)
        grid = torchvision.utils.make_grid(generated_images, nrow=num_rows)
        wandb.log({"Generated Images": [wandb.Image(grid, caption=f"Step {num}")]})

        # moving averages
        generated_images = self.generate_truncated(self.GAN.SE, self.GAN.GE, latents, n, trunc_psi=self.trunc_psi)
        grid = torchvision.utils.make_grid(generated_images, nrow=num_rows)
        wandb.log({"Generated Images EMA": [wandb.Image(grid, caption=f"Step {num}")]})

        # mixing regularities
        def tile(a, dim, n_tile):
            init_dim = a.size(dim)
            repeat_idx = [1] * a.dim()
            repeat_idx[dim] = n_tile
            a = a.repeat(*(repeat_idx))
            order_index = torch.LongTensor(
                np.concatenate([init_dim * np.arange(n_tile) + i for i in range(init_dim)])
            ).cuda()
            return torch.index_select(a, dim, order_index)

        nn = noise(num_rows, latent_dim)
        tmp1 = tile(nn, 0, num_rows)
        tmp2 = nn.repeat(num_rows, 1)

        tt = int(num_layers / 2)
        mixed_latents = [(tmp1, tt), (tmp2, num_layers - tt)]

        generated_images = self.generate_truncated(self.GAN.SE, self.GAN.GE, mixed_latents, n, trunc_psi=self.trunc_psi)
        grid = torchvision.utils.make_grid(generated_images, nrow=num_rows)
        wandb.log({"Style Mixing": [wandb.Image(grid, caption=f"Step {num}")]})

    @torch.no_grad()
    def calculate_fid(self):
        self.GAN.eval()

        inception = InceptionV3([3], normalize_input=False, init_weights=False)
        inception = inception.eval().to(next(self.GAN.parameters()).device)

        latent_dim = self.GAN.G.latent_dim
        image_size = self.GAN.G.image_size
        num_layers = self.GAN.G.num_layers

        features = []
        for _ in range(floor(self.fid_n_sample / self.batch_size) + 1):
            latents = noise_list(self.batch_size, num_layers, latent_dim)
            n = image_noise(self.batch_size, image_size)
            imgs = self.generate_truncated(self.GAN.SE, self.GAN.GE, latents, n, trunc_psi=self.trunc_psi)
            feat = inception(imgs)[0].view(imgs.shape[0], -1)
            features.append(feat.to("cpu"))
        features = torch.cat(features, 0).numpy()

        del inception

        sample_mean = np.mean(features, 0)
        sample_cov = np.cov(features, rowvar=False)

        with open(f"inception_{self.name}_stats.pkl", "rb") as f:
            embeds = pickle.load(f)
            real_mean = embeds["mean"]
            real_cov = embeds["cov"]

        cov_sqrt, _ = linalg.sqrtm(sample_cov @ real_cov, disp=False)
        if not np.isfinite(cov_sqrt).all():
            print("product of cov matrices is singular")
            offset = np.eye(sample_cov.shape[0]) * 1e-6
            cov_sqrt = linalg.sqrtm((sample_cov + offset) @ (real_cov + offset))
        if np.iscomplexobj(cov_sqrt):
            if not np.allclose(np.diagonal(cov_sqrt).imag, 0, atol=1e-3):
                raise ValueError(f"Imaginary component {np.max(np.abs(cov_sqrt.imag))}")
            cov_sqrt = cov_sqrt.real

        mean_diff = sample_mean - real_mean
        mean_norm = mean_diff @ mean_diff
        trace = np.trace(sample_cov) + np.trace(real_cov) - 2 * np.trace(cov_sqrt)
        inception_distance = mean_norm + trace

        with open(f"inception_{self.name}_features.pkl", "rb") as f:
            embeds = pickle.load(f)
            real_feats = embeds["features"]
        _, _, density, coverage = validation.prdc(real_feats[:80000], features[:80000])

        return inception_distance, density, coverage

    @torch.no_grad()
    def calculate_ppl(self):
        self.GAN.eval()
        latent_dim = self.GAN.G.latent_dim
        image_size = self.GAN.G.image_size
        num_layers = self.GAN.G.num_layers
        eps = 1e-4

        def lerp(a, b, t):
            return a + (b - a) * t

        percept = lpips.PerceptualLoss(
            model="net-lin", net="vgg", use_gpu=True, gpu_ids=[next(self.GAN.parameters()).device.index]
        )

        distances = []
        for _ in range(floor(self.fid_n_sample / self.batch_size) + 1):
            noise = image_noise(self.batch_size * 2, image_size)
            inputs = noise_list(self.batch_size * 2, num_layers, latent_dim)
            lerp_t = torch.rand(self.batch_size).cuda()

            # print(lerp_t.shape)

            w_space = []
            for tensor, num_layers in inputs:
                av = torch.from_numpy(self.av).cuda()
                tmp = self.trunc_psi * (self.GAN.SE(tensor) - av) + av
                w_space.append((tmp, num_layers))
            latent = styles_def_to_tensor(w_space)

            latent_t0, latent_t1 = latent[::2], latent[1::2]
            latent_e0 = lerp(latent_t0, latent_t1, lerp_t[:, None, None])
            latent_e1 = lerp(latent_t0, latent_t1, lerp_t[:, None, None] + eps)
            latent_e = torch.stack([latent_e0, latent_e1], 1).view(*latent.shape)

            image = evaluate_in_chunks(self.batch_size, self.GAN.GE, latent_e, noise)

            factor = image.shape[2] // 256
            if factor > 1:
                image = F.interpolate(image, size=(256, 256), mode="bilinear", align_corners=False)

            dist = percept(image[::2], image[1::2]).view(image.shape[0] // 2) / (eps ** 2)
            distances.append(dist.to("cpu").numpy())
        distances = np.concatenate(distances, 0)

        lo = np.percentile(distances, 1, interpolation="lower")
        hi = np.percentile(distances, 99, interpolation="higher")
        filtered_dist = np.extract(np.logical_and(lo <= distances, distances <= hi), distances)
        path_length = filtered_dist.mean()

        del percept, inputs, lerp_t, image, dist

        return path_length

    @torch.no_grad()
    def generate_truncated(self, S, G, style, noi, trunc_psi=0.75, num_image_tiles=8):
        latent_dim = G.latent_dim

        if self.av is None:
            z = noise(2 ** 14, latent_dim)
            samples = evaluate_in_chunks(self.batch_size, S, z).cpu().numpy()
            self.av = np.mean(samples, axis=0)
            self.av = np.expand_dims(self.av, axis=0)

        w_space = []
        for tensor, num_layers in style:
            tmp = S(tensor)
            av_torch = torch.from_numpy(self.av).cuda()
            tmp = trunc_psi * (tmp - av_torch) + av_torch
            w_space.append((tmp, num_layers))

        w_styles = styles_def_to_tensor(w_space)
        generated_images = evaluate_in_chunks(self.batch_size, G, w_styles, noi)
        return generated_images.clamp_(0.0, 1.0)

    def print_log(self):
        print(
            f"G: {self.g_loss:.2f} | D: {self.d_loss:.2f} | GP: {self.last_gp_loss:.2f} | PL: {self.pl_mean:.2f} | CR: {self.last_cr_loss:.2f} | Q: {self.q_loss:.2f}"
        )

    def model_name(self, num):
        return str(self.models_dir / self.name / f"model_{wandb.run.dir.split('/')[-1].split('-')[-1]}_{num}.pt")

    def init_folders(self):
        (self.results_dir / self.name).mkdir(parents=True, exist_ok=True)
        (self.models_dir / self.name).mkdir(parents=True, exist_ok=True)

    def clear(self):
        rmtree(f"./models/{self.name}", True)
        rmtree(f"./results/{self.name}", True)
        rmtree(str(self.config_path), True)
        self.init_folders()

    def save(self, num):
        torch.save(self.GAN.state_dict(), self.model_name(num))
        self.write_config()

    def load(self, num=-1):
        self.load_config()

        name = num
        if num == -1:
            file_paths = [p for p in Path(self.models_dir / self.name).glob("model_*.pt")]
            saved_nums = sorted(map(lambda x: int(x.stem.split("_")[-1]), file_paths))
            if len(saved_nums) == 0:
                return
            name = saved_nums[-1]
            print(f"continuing from previous epoch - {name}")
        self.steps = name * self.save_every
        self.GAN.load_state_dict(torch.load(self.model_name(name)))


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("data", type=str)
    parser.add_argument("name", type=str)
    parser.add_argument("--results_dir", type=str, default="/home/hans/neurout/")
    parser.add_argument("--models_dir", type=str, default="/home/hans/modelzoo/maua-sg2/")
    parser.add_argument("--new", type=bool, default=False)
    parser.add_argument("--load_from", type=str, default=-1)
    parser.add_argument("--image_size", type=int, default=256)
    parser.add_argument("--network_capacity", type=int, default=16)
    parser.add_argument("--transparent", type=bool, default=False)
    parser.add_argument("--batch_size", type=int, default=2)
    parser.add_argument("--gradient_accumulate_every", type=int, default=12)
    parser.add_argument("--num_train_steps", type=int, default=150000)
    parser.add_argument("--learning_rate", type=float, default=2e-4)
    parser.add_argument("--num_workers", type=int, default=None)
    parser.add_argument("--save_every", type=int, default=1000)
    parser.add_argument("--generate", type=bool, default=False)
    parser.add_argument("--num_image_tiles", type=int, default=8)
    parser.add_argument("--trunc_psi", type=float, default=1)
    parser.add_argument("--fp16", type=bool, default=False)
    parser.add_argument("--cl_reg", type=bool, default=True)
    parser.add_argument("--fq_layers", default=[])
    parser.add_argument("--fq_dict_size", default=256)
    parser.add_argument("--attn_layers", default=[])
    args = parser.parse_args()

    wandb.init(project=f"maua-stylegan", name="lucidrains-" + args.name)

    model = Trainer(
        args.name,
        args.results_dir,
        args.models_dir,
        batch_size=args.batch_size,
        gradient_accumulate_every=args.gradient_accumulate_every,
        image_size=args.image_size,
        network_capacity=args.network_capacity,
        transparent=args.transparent,
        lr=args.learning_rate,
        num_workers=args.num_workers,
        save_every=args.save_every,
        trunc_psi=args.trunc_psi,
        fp16=args.fp16,
        cl_reg=args.cl_reg,
        fq_layers=args.fq_layers,
        fq_dict_size=args.fq_dict_size,
        attn_layers=args.attn_layers,
    )

    if not args.new:
        model.load(args.load_from)
    else:
        model.clear()

    if args.generate:
        now = datetime.now()
        timestamp = now.strftime("%m-%d-%Y_%H-%M-%S")
        samples_name = f"generated-{timestamp}"
        model.evaluate(samples_name, args.num_image_tiles)
        print(f"sample images generated at {args.results_dir}/{args.name}/{args.samples_name}")
        exit()

    model.set_data_src(args.data)

    PBAR = tqdm(range(args.num_train_steps - model.steps), mininterval=10.0, desc=f"{args.name}<{args.data}>")
    for _ in PBAR:
        retry_call(model.train, tries=3, exceptions=NanException)
        if _ % 50 == 0:
            model.print_log()


================================================
FILE: models/autoencoder.py
================================================
import os
import sys
from copy import copy

import torch as th
import torch.nn.functional as F

sys.path.append(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))

from op import FusedLeakyReLU


def info(x):
    print(x.shape, x.min(), x.mean(), x.max())


class PrintShape(th.nn.Module):
    def __init__(self):
        super(PrintShape, self).__init__()

    def forward(self, x):
        print(x.shape)
        return x


class Flatten(th.nn.Module):
    def forward(self, x):
        return x.view(x.size(0), -1)


class UnFlatten(th.nn.Module):
    def __init__(self, channels, size):
        super(UnFlatten, self).__init__()
        self.channels = channels
        self.size = size

    def forward(self, x):
        return x.view(x.size(0), self.channels, self.size, self.size)


class LogCoshVAE(th.nn.Module):
    """
    Adapted from https://github.com/AntixK/PyTorch-VAE
    See LICENSE_AUTOENCODER
    """

    def __init__(self, in_channels, latent_dim, hidden_dims=None, alpha=10.0, beta=1.0, kld_weight=1):
        super(LogCoshVAE, self).__init__()

        my_hidden_dims = copy(hidden_dims)

        self.latent_dim = latent_dim
        self.alpha = alpha
        self.beta = beta
        self.kld_weight = kld_weight

        modules = []
        if my_hidden_dims is None:
            my_hidden_dims = [32, 64, 128, 256, 512]

        # Build Encoder
        for h_dim in my_hidden_dims:
            modules.append(
                th.nn.Sequential(
                    th.nn.Conv2d(in_channels, out_channels=h_dim, kernel_size=3, stride=2, padding=1),
                    th.nn.BatchNorm2d(h_dim),
                    FusedLeakyReLU(h_dim),
                )
            )
            in_channels = h_dim

        self.encoder = th.nn.Sequential(*modules)
        self.fc_mu = th.nn.Linear(my_hidden_dims[-1] * 4, latent_dim)
        self.fc_var = th.nn.Linear(my_hidden_dims[-1] * 4, latent_dim)

        # Build Decoder
        modules = []

        self.decoder_input = th.nn.Linear(latent_dim, my_hidden_dims[-1] * 4)

        my_hidden_dims.reverse()

        for i in range(len(my_hidden_dims) - 1):
            modules.append(
                th.nn.Sequential(
                    th.nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False),
                    th.nn.Conv2d(my_hidden_dims[i], my_hidden_dims[i + 1], kernel_size=3, padding=1),
                    th.nn.BatchNorm2d(my_hidden_dims[i + 1]),
                    FusedLeakyReLU(my_hidden_dims[i + 1]),
                )
            )

        self.decoder = th.nn.Sequential(*modules)

        self.final_layer = th.nn.Sequential(
            th.nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False),
            th.nn.Conv2d(my_hidden_dims[-1], my_hidden_dims[-1], kernel_size=3, padding=1),
            th.nn.BatchNorm2d(my_hidden_dims[-1]),
            FusedLeakyReLU(my_hidden_dims[-1]),
            th.nn.Conv2d(my_hidden_dims[-1], out_channels=3, kernel_size=3, padding=1),
            th.nn.Tanh(),
        )

    def encode(self, input):
        result = self.encoder(input)
        result = th.flatten(result, start_dim=1)

        mu = self.fc_mu(result)
        log_var = self.fc_var(result)

        return mu, log_var

    def decode(self, z):
        result = self.decoder_input(z)
        result = result.view(-1, self.latent_dim, 2, 2)
        result = self.decoder(result)
        result = self.final_layer(result)
        return result

    def reparameterize(self, mu, logvar):
        std = th.exp(0.5 * logvar)
        eps = th.randn_like(std)
        return eps * std + mu

    def forward(self, input):
        mu, log_var = self.encode(input)
        z = self.reparameterize(mu, log_var)
        return self.decode(z), mu, log_var

    def loss(self, real, fake, mu, log_var):
        t = fake - real

        recons_loss = self.alpha * t + th.log(1.0 + th.exp(-2 * self.alpha * t)) - th.log(2.0 * th.ones((1)))
        recons_loss = (1.0 / self.alpha) * recons_loss.mean()

        kld_loss = th.mean(-0.5 * th.sum(1 + log_var - mu ** 2 - log_var.exp(), dim=1), dim=0)

        loss = recons_loss + self.beta * self.kld_weight * kld_loss

        return {"Total": loss, "Reconstruction": recons_loss, "Kullback Leibler Divergence": -kld_loss}


class conv2DBatchNormRelu(th.nn.Module):
    def __init__(
        self, in_channels, n_filters, k_size, stride, padding, bias=True, dilation=1, with_bn=True,
    ):
        super(conv2DBatchNormRelu, self).__init__()

        conv_mod = th.nn.Conv2d(
            int(in_channels),
            int(n_filters),
            kernel_size=k_size,
            padding=padding,
            stride=stride,
            bias=bias,
            dilation=dilation,
        )

        if with_bn:
            self.cbr_unit = th.nn.Sequential(
                conv_mod, th.nn.BatchNorm2d(int(n_filters)), FusedLeakyReLU(int(n_filters))
            )
        else:
            self.cbr_unit = th.nn.Sequential(conv_mod, FusedLeakyReLU(int(n_filters)))

    def forward(self, inputs):
        outputs = self.cbr_unit(inputs)
        return outputs


class segnetDown2(th.nn.Module):
    def __init__(self, in_size, out_size):
        super(segnetDown2, self).__init__()
        self.conv1 = conv2DBatchNormRelu(in_size, out_size, 3, 1, 1)
        self.conv2 = conv2DBatchNormRelu(out_size, out_size, 3, 1, 1)
        self.maxpool_with_argmax = th.nn.MaxPool2d(2, 2, return_indices=True)

    def forward(self, inputs):
        outputs = self.conv1(inputs)
        outputs = self.conv2(outputs)
        unpooled_shape = outputs.size()
        outputs, indices = self.maxpool_with_argmax(outputs)
        return outputs, indices, unpooled_shape


class segnetDown3(th.nn.Module):
    def __init__(self, in_size, out_size):
        super(segnetDown3, self).__init__()
        self.conv1 = conv2DBatchNormRelu(in_size, out_size, 3, 1, 1)
        self.conv2 = conv2DBatchNormRelu(out_size, out_size, 3, 1, 1)
        self.conv3 = conv2DBatchNormRelu(out_size, out_size, 3, 1, 1)
        self.maxpool_with_argmax = th.nn.MaxPool2d(2, 2, return_indices=True)

    def forward(self, inputs):
        outputs = self.conv1(inputs)
        outputs = self.conv2(outputs)
        outputs = self.conv3(outputs)
        unpooled_shape = outputs.size()
        outputs, indices = self.maxpool_with_argmax(outputs)
        return outputs, indices, unpooled_shape


class segnetUp2(th.nn.Module):
    def __init__(self, in_size, out_size):
        super(segnetUp2, self).__init__()
        self.unpool = th.nn.MaxUnpool2d(2, 2)
        self.conv1 = conv2DBatchNormRelu(in_size, in_size, 3, 1, 1)
        self.conv2 = conv2DBatchNormRelu(in_size, out_size, 3, 1, 1)

    def forward(self, inputs, indices, output_shape):
        outputs = self.unpool(input=inputs, indices=indices, output_size=output_shape)
        outputs = self.conv1(outputs)
        outputs = self.conv2(outputs)
        return outputs


class segnetUp3(th.nn.Module):
    def __init__(self, in_size, out_size):
        super(segnetUp3, self).__init__()
        self.unpool = th.nn.MaxUnpool2d(2, 2)
        self.conv1 = conv2DBatchNormRelu(in_size, in_size, 3, 1, 1)
        self.conv2 = conv2DBatchNormRelu(in_size, in_size, 3, 1, 1)
        self.conv3 = conv2DBatchNormRelu(in_size, out_size, 3, 1, 1)

    def forward(self, inputs, indices, output_shape):
        outputs = self.unpool(input=inputs, indices=indices, output_size=output_shape)
        outputs = self.conv1(outputs)
        outputs = self.conv2(outputs)
        outputs = self.conv3(outputs)
        return outputs


class SegNet(th.nn.Module):
    """
    Adapted from https://github.com/foamliu/Autoencoder
    See LICENSE_AUTOENCODER
    """

    def __init__(self, in_channels=3):
        super(SegNet, self).__init__()

        self.down1 = segnetDown2(in_channels, 64)
        self.down2 = segnetDown2(64, 128)
        self.down3 = segnetDown3(128, 256)
        self.down4 = segnetDown3(256, 512)
        self.down5 = segnetDown3(512, 512)

        self.up5 = segnetUp3(512, 512)
        self.up4 = segnetUp3(512, 256)
        self.up3 = segnetUp3(256, 128)
        self.up2 = segnetUp2(128, 64)
        self.up1 = segnetUp2(64, in_channels)

    def random_indices(self, shape):
        batch, channel, height, width = shape
        xy = th.randint(0, 2, size=[batch, channel, height, width, 2])
        grid = th.arange(height * width).reshape(height, width)
        indices = grid * 2 + (th.arange(height) * width * 2)[:, None] + xy[..., 0] + width * 2 * xy[..., 1]
        return indices.cuda()

    def encode(self, inputs):
        down1, indices_1, unpool_shape1 = self.down1(inputs)
        down2, indices_2, unpool_shape2 = self.down2(down1)
        down3, indices_3, unpool_shape3 = self.down3(down2)
        down4, indices_4, unpool_shape4 = self.down4(down3)
        down5, indices_5, unpool_shape5 = self.down5(down4)
        return down5

    def decode(self, inp):
        batch, _, height, width = inp.shape
        up5 = self.up5(inp, self.random_indices([batch, 512, height, width]), [batch, 512, height * 2, width * 2])
        up4 = self.up4(
            up5, self.random_indices([batch, 512, height * 2, width * 2]), [batch, 512, height * 4, width * 4]
        )
        up3 = self.up3(
            up4, self.random_indices([batch, 256, height * 4, width * 4]), [batch, 256, height * 8, width * 8]
        )
        up2 = self.up2(
            up3, self.random_indices([batch, 128, height * 8, width * 8]), [batch, 128, height * 16, width * 16]
        )
        up1 = self.up1(
            up2, self.random_indices([batch, 64, height * 16, width * 16]), [batch, 64, height * 32, width * 32]
        )
        return up1

    def forward(self, inputs):
        down1, indices_1, unpool_shape1 = self.down1(inputs)
        down2, indices_2, unpool_shape2 = self.down2(down1)
        down3, indices_3, unpool_shape3 = self.down3(down2)
        down4, indices_4, unpool_shape4 = self.down4(down3)
        down5, indices_5, unpool_shape5 = self.down5(down4)

        up5 = self.up5(down5, indices_5.shape, unpool_shape5)
        up4 = self.up4(up5, indices_4.shape, unpool_shape4)
        up3 = self.up3(up4, indices_3.shape, unpool_shape3)
        up2 = self.up2(up3, indices_2.shape, unpool_shape2)
        up1 = self.up1(up2, indices_1.shape, unpool_shape1)

        return up1

    def init_vgg16_params(self, vgg16):
        blocks = [self.down1, self.down2, self.down3, self.down4, self.down5]

        ranges = [[0, 4], [5, 9], [10, 16], [17, 23], [24, 29]]
        features = list(vgg16.features.children())

        vgg_layers = []
        for _layer in features:
            if isinstance(_layer, th.nn.Conv2d):
                vgg_layers.append(_layer)

        merged_layers = []
        for idx, conv_block in enumerate(blocks):
            if idx < 2:
                units = [conv_block.conv1.cbr_unit, conv_block.conv2.cbr_unit]
            else:
                units = [
                    conv_block.conv1.cbr_unit,
                    conv_block.conv2.cbr_unit,
                    conv_block.conv3.cbr_unit,
                ]
            for _unit in units:
                for _layer in _unit:
                    if isinstance(_layer, th.nn.Conv2d):
                        merged_layers.append(_layer)

        assert len(vgg_layers) == len(merged_layers)

        for l1, l2 in zip(vgg_layers, merged_layers):
            if isinstance(l1, th.nn.Conv2d) and isinstance(l2, th.nn.Conv2d):
                assert l1.weight.size() == l2.weight.size()
                assert l1.bias.size() == l2.bias.size()
                l2.weight.data = l1.weight.data
                l2.bias.data = l1.bias.data


class ConvSegNet(th.nn.Module):
    """
    Adapted from https://github.com/foamliu/Autoencoder
    See LICENSE_AUTOENCODER
    """

    def __init__(self, in_channels=3):
        super(ConvSegNet, self).__init__()

        self.encoder = th.nn.Sequential(
            conv2DBatchNormRelu(in_channels, 64, 3, 1, 1),
            conv2DBatchNormRelu(64, 64, 3, 1, 1),
            th.nn.MaxPool2d(2, 2),
            conv2DBatchNormRelu(64, 128, 3, 1, 1),
            conv2DBatchNormRelu(128, 128, 3, 1, 1),
            th.nn.MaxPool2d(2, 2),
            conv2DBatchNormRelu(128, 256, 3, 1, 1),
            conv2DBatchNormRelu(256, 256, 3, 1, 1),
            conv2DBatchNormRelu(256, 256, 3, 1, 1),
            th.nn.MaxPool2d(2, 2),
            conv2DBatchNormRelu(256, 512, 3, 1, 1),
            conv2DBatchNormRelu(512, 512, 3, 1, 1),
            conv2DBatchNormRelu(512, 512, 3, 1, 1),
            th.nn.MaxPool2d(2, 2),
            conv2DBatchNormRelu(512, 512, 3, 1, 1),
            conv2DBatchNormRelu(512, 512, 3, 1, 1),
            conv2DBatchNormRelu(512, 512, 3, 1, 1),
            th.nn.MaxPool2d(2, 2),
            th.nn.Tanh(),
        )

        self.decoder = th.nn.Sequential(
            th.nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False),
            conv2DBatchNormRelu(512, 512, 3, 1, 1),
            conv2DBatchNormRelu(512, 512, 3, 1, 1),
            conv2DBatchNormRelu(512, 512, 3, 1, 1),
            th.nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False),
            conv2DBatchNormRelu(512, 512, 3, 1, 1),
            conv2DBatchNormRelu(512, 512, 3, 1, 1),
            conv2DBatchNormRelu(512, 256, 3, 1, 1),
            th.nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False),
            conv2DBatchNormRelu(256, 256, 3, 1, 1),
            conv2DBatchNormRelu(256, 256, 3, 1, 1),
            conv2DBatchNormRelu(256, 128, 3, 1, 1),
            th.nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False),
            conv2DBatchNormRelu(128, 128, 3, 1, 1),
            conv2DBatchNormRelu(128, 64, 3, 1, 1),
            th.nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False),
            conv2DBatchNormRelu(64, 64, 3, 1, 1),
            conv2DBatchNormRelu(64, in_channels, 3, 1, 1),
        )

    def encode(self, inputs):
        return self.encoder(inputs)

    def decode(self, inputs):
        return self.decoder(inputs)

    def forward(self, inputs):
        z = self.encode(inputs)
        # print(z.min(), z.mean(), z.max(), z.shape)
        return self.decode(z)


class VariationalConvSegNet(th.nn.Module):
    """
    Adapted from https://github.com/foamliu/Autoencoder
    See LICENSE_AUTOENCODER
    """

    def __init__(self, in_channels=3):
        super(VariationalConvSegNet, self).__init__()

        self.encoder = th.nn.Sequential(
            conv2DBatchNormRelu(in_channels, 64, 3, 1, 1),
            conv2DBatchNormRelu(64, 64, 3, 1, 1),
            th.nn.MaxPool2d(2, 2),
            conv2DBatchNormRelu(64, 128, 3, 1, 1),
            conv2DBatchNormRelu(128, 128, 3, 1, 1),
            th.nn.MaxPool2d(2, 2),
            conv2DBatchNormRelu(128, 256, 3, 1, 1),
            conv2DBatchNormRelu(256, 256, 3, 1, 1),
            conv2DBatchNormRelu(256, 256, 3, 1, 1),
            th.nn.MaxPool2d(2, 2),
            conv2DBatchNormRelu(256, 512, 3, 1, 1),
            conv2DBatchNormRelu(512, 512, 3, 1, 1),
            conv2DBatchNormRelu(512, 512, 3, 1, 1),
            th.nn.MaxPool2d(2, 2),
            conv2DBatchNormRelu(512, 512, 3, 1, 1),
            conv2DBatchNormRelu(512, 512, 3, 1, 1),
            conv2DBatchNormRelu(512, 512, 3, 1, 1),
            th.nn.MaxPool2d(2, 2),
            th.nn.Tanh(),
            Flatten(),
        )

        self.fc_mu = th.nn.Linear(512 * 4 * 4, 512 * 4 * 4)
        self.fc_var = th.nn.Linear(512 * 4 * 4, 512 * 4 * 4)

        self.decoder = th.nn.Sequential(
            UnFlatten(512, 4),
            th.nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False),
            conv2DBatchNormRelu(512, 512, 3, 1, 1),
            conv2DBatchNormRelu(512, 512, 3, 1, 1),
            conv2DBatchNormRelu(512, 512, 3, 1, 1),
            th.nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False),
            conv2DBatchNormRelu(512, 512, 3, 1, 1),
            conv2DBatchNormRelu(512, 512, 3, 1, 1),
            conv2DBatchNormRelu(512, 256, 3, 1, 1),
            th.nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False),
            conv2DBatchNormRelu(256, 256, 3, 1, 1),
            conv2DBatchNormRelu(256, 256, 3, 1, 1),
            conv2DBatchNormRelu(256, 128, 3, 1, 1),
            th.nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False),
            conv2DBatchNormRelu(128, 128, 3, 1, 1),
            conv2DBatchNormRelu(128, 64, 3, 1, 1),
            th.nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False),
            conv2DBatchNormRelu(64, 64, 3, 1, 1),
            conv2DBatchNormRelu(64, in_channels, 3, 1, 1),
            th.nn.Tanh(),
        )

    def reparameterize(self, mu, log_var):
        std = th.exp(0.5 * log_var)
        eps = th.randn_like(std)
        return eps * std + mu

    def encode(self, inputs):
        result = self.encoder(inputs)

        mu = self.fc_mu(result)
        log_var = self.fc_var(result)

        return mu, log_var

    def decode(self, inputs):
        return self.decoder(inputs)

    def forward(self, inputs):
        mu, log_var = self.encode(inputs)
        z = self.reparameterize(mu, log_var)
        return self.decode(z)


def create_encoder_single_conv(in_chs, out_chs, kernel):
    assert kernel % 2 == 1
    return th.nn.Sequential(
        th.nn.Conv2d(in_chs, out_chs, kernel_size=kernel, padding=(kernel - 1) // 2),
        th.nn.BatchNorm2d(out_chs),
        FusedLeakyReLU(out_chs),
    )


class EncoderInceptionModuleSignle(th.nn.Module):
    def __init__(self, channels):
        assert channels % 2 == 0
        super().__init__()
        # put bottle-neck layers before convolution
        bn_ch = channels // 2
        self.bottleneck = create_encoder_single_conv(channels, bn_ch, 1)
        # bn -> Conv1, 3, 5
        self.conv1 = create_encoder_single_conv(bn_ch, channels, 1)
        self.conv3 = create_encoder_single_conv(bn_ch, channels, 3)
        self.conv5 = create_encoder_single_conv(bn_ch, channels, 5)
        self.conv7 = create_encoder_single_conv(bn_ch, channels, 7)
        # pool-proj(no-bottle neck)
        self.pool3 = th.nn.MaxPool2d(3, stride=1, padding=1)
        self.pool5 = th.nn.MaxPool2d(5, stride=1, padding=2)

    def forward(self, x):
        # Original inception is concatenation, but use simple addition instead
        bn = self.bottleneck(x)
        out = self.conv1(bn) + self.conv3(bn) + self.conv5(bn) + self.conv7(bn) + self.pool3(x) + self.pool5(x)
        return out


class EncoderModule(th.nn.Module):
    def __init__(self, chs, repeat_num, use_inception):
        super().__init__()
        if use_inception:
            layers = [EncoderInceptionModuleSignle(chs) for i in range(repeat_num)]
        else:
            layers = [create_encoder_single_conv(chs, chs, 3) for i in range(repeat_num)]
        self.convs = th.nn.Sequential(*layers)

    def forward(self, x):
        return self.convs(x)


class Encoder(th.nn.Module):
    def __init__(self, use_inception, repeat_per_module):
        super().__init__()
        # stages
        self.upch1 = th.nn.Conv2d(3, 32, kernel_size=3)
        self.stage1 = EncoderModule(32, repeat_per_module, use_inception)
        self.upch2 = self._create_downsampling_module(32, 2)
        self.stage2 = EncoderModule(64, repeat_per_module, use_inception)
        self.upch3 = self._create_downsampling_module(64, 2)
        self.stage3 = EncoderModule(128, repeat_per_module, use_inception)
        self.upch4 = self._create_downsampling_module(128, 2)
        self.stage4 = EncoderModule(256, repeat_per_module, use_inception)

    def _create_downsampling_module(self, input_channels, pooling_kenel):
        return th.nn.Sequential(
            th.nn.AvgPool2d(pooling_kenel),
            th.nn.Conv2d(input_channels, input_channels * 2, kernel_size=1),
            th.nn.BatchNorm2d(input_channels * 2),
            FusedLeakyReLU(input_channels * 2),
        )

    def forward(self, x):
        # print(x.shape)
        out = self.stage1(self.upch1(x))
        # print(out.shape)
        out = self.stage2(self.upch2(out))
        # print(out.shape)
        out = self.stage3(self.upch3(out))
        # print(out.shape)
        out = self.stage4(self.upch4(out))
        # print(out.shape)
        out = F.avg_pool2d(out, 8)  # Global Average pooling
        # print(out.shape)
        return out.view(-1, 256)


## Decoder
def create_decoder_single_conv(in_chs, out_chs, kernel):
    assert kernel % 2 == 1
    return th.nn.Sequential(
        th.nn.ConvTranspose2d(in_chs, out_chs, kernel_size=kernel, padding=(kernel - 1) // 2),
        th.nn.BatchNorm2d(out_chs),
        FusedLeakyReLU(out_chs),
    )


class DecoderInceptionModuleSingle(th.nn.Module):
    def __init__(self, channels):
        assert channels % 2 == 0
        super().__init__()
        # put bottle-neck layers before convolution
        bn_ch = channels // 4
        self.bottleneck = create_decoder_single_conv(channels, bn_ch, 1)
        # bn -> Conv1, 3, 5
        self.conv1 = create_decoder_single_conv(bn_ch, channels, 1)
        self.conv3 = create_decoder_single_conv(bn_ch, channels, 3)
        self.conv5 = create_decoder_single_conv(bn_ch, channels, 5)
        self.conv7 = create_decoder_single_conv(bn_ch, channels, 7)
        # pool-proj(no-bottle neck)
        self.pool3 = th.nn.MaxPool2d(3, stride=1, padding=1)
        self.pool5 = th.nn.MaxPool2d(5, stride=1, padding=2)

    def forward(self, x):
        # Original inception is concatenation, but use simple addition instead
        bn = self.bottleneck(x)
        out = self.conv1(bn) + self.conv3(bn) + self.conv5(bn) + self.conv7(bn) + self.pool3(x) + self.pool5(x)
        return out


class DecoderModule(th.nn.Module):
    def __init__(self, chs, repeat_num, use_inception):
        super().__init__()
        if use_inception:
            layers = [DecoderInceptionModuleSingle(chs) for i in range(repeat_num)]
        else:
            layers = [create_decoder_single_conv(chs, chs, 3) for i in range(repeat_num)]
        self.convs = th.nn.Sequential(*layers)

    def forward(self, x):
        return self.convs(x)


class Decoder(th.nn.Module):
    def __init__(self, use_inception, repeat_per_module):
        super().__init__()
        # stages
        self.stage1 = DecoderModule(256, repeat_per_module, use_inception)
        self.downch1 = self._create_upsampling_module(256, 2)
        self.stage2 = DecoderModule(128, repeat_per_module, use_inception)
        self.downch2 = self._create_upsampling_module(128, 2)
        self.stage3 = DecoderModule(64, repeat_per_module, use_inception)
        self.downch3 = self._create_upsampling_module(64, 2)
        self.stage4 = DecoderModule(32, repeat_per_module, use_inception)
        self.downch4 = self._create_upsampling_module(32, 2)
        self.last = th.nn.ConvTranspose2d(16, 3, kernel_size=1)

    def _create_upsampling_module(self, input_channels, pooling_kenel):
        return th.nn.Sequential(
            th.nn.ConvTranspose2d(input_channels, input_channels // 2, kernel_size=pooling_kenel, stride=pooling_kenel),
            th.nn.BatchNorm2d(input_channels // 2),
            FusedLeakyReLU(input_channels // 2),
        )

    def forward(self, x):
        out = F.upsample(x.view(-1, 256, 1, 1), scale_factor=8)
        out = self.downch1(self.stage1(out))
        out = self.downch2(self.stage2(out))
        out = self.downch3(self.stage3(out))
        out = self.downch4(self.stage4(out))
        return th.sigmoid(self.last(out))


## VAE
class InceptionVAE(th.nn.Module):
    """
    Adapted from https://github.com/koshian2/inception-vae
    """

    def __init__(self, latent_dim=512, repeat_per_block=1, use_inception=True):
        super(InceptionVAE, self).__init__()

        # # latent features
        self.n_latent_features = latent_dim

        # Encoder
        self.encoder = Encoder(use_inception, repeat_per_block)
        # Middle
        self.fc_mu = th.nn.Linear(256, self.n_latent_features)
        self.fc_logvar = th.nn.Linear(256, self.n_latent_features)
        self.fc_rep = th.nn.Linear(self.n_latent_features, 256)
        # Decoder
        self.decoder = Decoder(use_inception, repeat_per_block)

    def _reparameterize(self, mu, logvar):
        std = logvar.mul(0.5).exp_()
        esp = th.randn(*mu.size()).cuda()
        z = mu + std * esp
        return z

    def _bottleneck(self, h):
        mu, logvar = self.fc_mu(h), self.fc_logvar(h)
        z = self._reparameterize(mu, logvar)
        return z, mu, logvar

    def sampling(self):
        # assume latent features space ~ N(0, 1)
        z = th.randn(24, self.n_latent_features).cuda()
        z = self.fc_rep(z)
        # decode
        return self.decoder(z)

    def forward(self, x):
        # Encoder
        h = self.encoder(x)
        # Bottle-neck
        z, mu, logvar = self._bottleneck(h)
        # decoder
        z = self.fc_rep(z)
        d = self.decoder(z)
        return d, mu, logvar


================================================
FILE: models/stylegan1.py
================================================
# from https://github.com/lernapparat/lernapparat/blob/master/style_gan/pyth_style_gan.ipynb

import gc
from collections import OrderedDict

import numpy as np
import torch as th
import torch.nn as nn
import torch.nn.functional as F


class MyLinear(nn.Module):
    """Linear layer with equalized learning rate and custom learning rate multiplier."""

    def __init__(
        self, input_size, output_size, gain=2 ** (0.5), use_wscale=False, lrmul=1, bias=True,
    ):
        super().__init__()
        he_std = gain * input_size ** (-0.5)  # He init
        # Equalized learning rate and custom learning rate multiplier.
        if use_wscale:
            init_std = 1.0 / lrmul
            self.w_mul = he_std * lrmul
        else:
            init_std = he_std / lrmul
            self.w_mul = lrmul
        self.weight = th.nn.Parameter(th.randn(output_size, input_size) * init_std)
        if bias:
            self.bias = th.nn.Parameter(th.zeros(output_size))
            self.b_mul = lrmul
        else:
            self.bias = None

    def forward(self, x):
        bias = self.bias
        if bias is not None:
            bias = bias * self.b_mul
        return F.linear(x, self.weight * self.w_mul, bias)


class MyConv2d(nn.Module):
    """Conv layer with equalized learning rate and custom learning rate multiplier."""

    def __init__(
        self,
        input_channels,
        output_channels,
        kernel_size,
        gain=2 ** (0.5),
        use_wscale=False,
        lrmul=1,
        bias=True,
        intermediate=None,
        upscale=False,
    ):
        super().__init__()
        if upscale:
            self.upscale = Upscale2d()
        else:
            self.upscale = None
        he_std = gain * (input_channels * kernel_size ** 2) ** (-0.5)  # He init
        self.kernel_size = kernel_size
        if use_wscale:
            init_std = 1.0 / lrmul
            self.w_mul = he_std * lrmul
        else:
            init_std = he_std / lrmul
            self.w_mul = lrmul
        self.weight = th.nn.Parameter(th.randn(output_channels, input_channels, kernel_size, kernel_size) * init_std)
        if bias:
            self.bias = th.nn.Parameter(th.zeros(output_channels))
            self.b_mul = lrmul
        else:
            self.bias = None
        self.intermediate = intermediate

    def forward(self, x):
        bias = self.bias
        if bias is not None:
            bias = bias * self.b_mul

        have_convolution = False
        if self.upscale is not None and min(x.shape[2:]) * 2 >= 128:
            # this is the fused upscale + conv from StyleGAN, sadly this seems incompatible with the non-fused way
            # this really needs to be cleaned up and go into the conv...
            w = self.weight * self.w_mul
            w = w.permute(1, 0, 2, 3)
            # probably applying a conv on w would be more efficient. also this quadruples the weight (average)?!
            w = F.pad(w, (1, 1, 1, 1))
            w = w[:, :, 1:, 1:] + w[:, :, :-1, 1:] + w[:, :, 1:, :-1] + w[:, :, :-1, :-1]
            x = F.conv_transpose2d(x, w, stride=2, padding=(w.size(-1) - 1) // 2)
            have_convolution = True
        elif self.upscale is not None:
            x = self.upscale(x)

        if not have_convolution and self.intermediate is None:
            return F.conv2d(x, self.weight * self.w_mul, bias, padding=self.kernel_size // 2)
        elif not have_convolution:
            x = F.conv2d(x, self.weight * self.w_mul, None, padding=self.kernel_size // 2)

        if self.intermediate is not None:
            x = self.intermediate(x)
        if bias is not None:
            x = x + bias.view(1, -1, 1, 1)
        return x


class NoiseLayer(nn.Module):
    """adds noise. noise is per pixel (constant over channels) with per-channel weight"""

    def __init__(self, channels):
        super().__init__()
        self.weight = nn.Parameter(th.zeros(channels))
        self.noise = None

    def forward(self, x):
        if self.noise is None:
            noise = th.randn(x.size(0), 1, x.size(2), x.size(3), device=x.device, dtype=x.dtype)
        else:
            noise = self.noise.to(x.device)
        # print(noise.shape, noise.min(), noise.mean(), noise.max())
        x = x + self.weight.view(1, -1, 1, 1) * noise
        return x


class StyleMod(nn.Module):
    def __init__(self, latent_size, channels, use_wscale):
        super(StyleMod, self).__init__()
        self.lin = MyLinear(latent_size, channels * 2, gain=1.0, use_wscale=use_wscale)

    def forward(self, x, latent):
        style = self.lin(latent)  # style => [batch_size, n_channels*2]
        shape = [-1, 2, x.size(1)] + (x.dim() - 2) * [1]
        style = style.view(shape)  # [batch_size, 2, n_channels, ...]
        x = x * (style[:, 0] + 1.0) + style[:, 1]
        return x


class PixelNormLayer(nn.Module):
    def __init__(self, epsilon=1e-8):
        super().__init__()
        self.epsilon = epsilon

    def forward(self, x):
        return x * th.rsqrt(th.mean(x ** 2, dim=1, keepdim=True) + self.epsilon)


class BlurLayer(nn.Module):
    def __init__(self, kernel=[1, 2, 1], normalize=True, flip=False, stride=1):
        super(BlurLayer, self).__init__()
        kernel = [1, 2, 1]
        kernel = th.tensor(kernel, dtype=th.float32)
        kernel = kernel[:, None] * kernel[None, :]
        kernel = kernel[None, None]
        if normalize:
            kernel = kernel / kernel.sum()
        if flip:
            kernel = kernel[:, :, ::-1, ::-1]
        self.register_buffer("kernel", kernel)
        self.stride = stride

    def forward(self, x):
        # expand kernel channels
        kernel = self.kernel.expand(x.size(1), -1, -1, -1)
        x = F.conv2d(x, kernel, stride=self.stride, padding=int((self.kernel.size(2) - 1) / 2), groups=x.size(1),)
        return x


def upscale2d(x, factor=2, gain=1):
    assert x.dim() == 4
    if gain != 1:
        x = x * gain
    if factor != 1:
        shape = x.shape
        x = x.view(shape[0], shape[1], shape[2], 1, shape[3], 1).expand(-1, -1, -1, factor, -1, factor)
        x = x.contiguous().view(shape[0], shape[1], factor * shape[2], factor * shape[3])
    return x


class Upscale2d(nn.Module):
    def __init__(self, factor=2, gain=1):
        super().__init__()
        assert isinstance(factor, int) and factor >= 1
        self.gain = gain
        self.factor = factor

    def forward(self, x):
        return upscale2d(x, factor=self.factor, gain=self.gain)


class G_mapping(nn.Sequential):
    def __init__(self, nonlinearity="lrelu", use_wscale=True):
        act, gain = {"relu": (th.relu, np.sqrt(2)), "lrelu": (nn.LeakyReLU(negative_slope=0.2), np.sqrt(2))}[
            nonlinearity
        ]
        layers = [
            ("pixel_norm", PixelNormLayer()),
            ("dense0", MyLinear(512, 512, gain=gain, lrmul=0.01, use_wscale=use_wscale),),
            ("dense0_act", act),
            ("dense1", MyLinear(512, 512, gain=gain, lrmul=0.01, use_wscale=use_wscale),),
            ("dense1_act", act),
            ("dense2", MyLinear(512, 512, gain=gain, lrmul=0.01, use_wscale=use_wscale),),
            ("dense2_act", act),
            ("dense3", MyLinear(512, 512, gain=gain, lrmul=0.01, use_wscale=use_wscale),),
            ("dense3_act", act),
            ("dense4", MyLinear(512, 512, gain=gain, lrmul=0.01, use_wscale=use_wscale),),
            ("dense4_act", act),
            ("dense5", MyLinear(512, 512, gain=gain, lrmul=0.01, use_wscale=use_wscale),),
            ("dense5_act", act),
            ("dense6", MyLinear(512, 512, gain=gain, lrmul=0.01, use_wscale=use_wscale),),
            ("dense6_act", act),
            ("dense7", MyLinear(512, 512, gain=gain, lrmul=0.01, use_wscale=use_wscale),),
            ("dense7_act", act),
        ]
        super().__init__(OrderedDict(layers))

    def forward(self, x):
        x = super().forward(x)
        # Broadcast
        x = x.unsqueeze(1).expand(-1, 18, -1)
        return x


class Truncation(nn.Module):
    def __init__(self, avg_latent, max_layer=8, threshold=0.7):
        super().__init__()
        self.max_layer = max_layer
        self.threshold = threshold
        self.register_buffer("avg_latent", avg_latent)

    def forward(self, x):
        assert x.dim() == 3
        interp = th.lerp(self.avg_latent, x, self.threshold)
        do_trunc = (th.arange(x.size(1)) < self.max_layer).view(1, -1, 1)
        return th.where(do_trunc, interp, x)


class LayerEpilogue(nn.Module):
    """Things to do at the end of each layer."""

    def __init__(
        self,
        channels,
        dlatent_size,
        use_wscale,
        use_noise,
        use_pixel_norm,
        use_instance_norm,
        use_styles,
        activation_layer,
    ):
        super().__init__()
        layers = []
        if use_noise:
            layers.append(("noise", NoiseLayer(channels)))
        layers.append(("activation", activation_layer))
        if use_pixel_norm:
            layers.append(("pixel_norm", PixelNormLayer()))
        if use_instance_norm:
            layers.append(("instance_norm", nn.InstanceNorm2d(channels)))
        self.top_epi = nn.Sequential(OrderedDict(layers))
        if use_styles:
            self.style_mod = StyleMod(dlatent_size, channels, use_wscale=use_wscale)
        else:
            self.style_mod = None

        # if use_noise:
        #     # layers.append(("noise", NoiseLayer(channels)))
        #     self.noise = NoiseLayer(channels)
        # else:
        #     self.noise = None

        # self.activation = activation_layer

        # if use_pixel_norm:
        #     self.pixel_norm = PixelNormLayer()
        # else:
        #     self.pixel_norm = None

        # if use_instance_norm:
        #     self.instance_norm = nn.InstanceNorm2d(channels)
        # else:
        #     self.instance_norm = None

        # if use_styles:
        #     self.style_mod = StyleMod(dlatent_size, channels, use_wscale=use_wscale)
        # else:
        #     self.style_mod = None

    def forward(self, x, dlatents_in_slice, noise):
        # if self.noise is not None:
        #     x = self.noise(x, noise)

        # x = self.activation(x)

        # if self.pixel_norm is not None:
        #     x = self.pixel_norm(x)

        # if self.instance_norm is not None:
        #     x = self.instance_norm(x)

        if noise is not None:
            self.top_epi.noise.noise = noise

        x = self.top_epi(x)
        if self.style_mod is not None:
            x = self.style_mod(x, dlatents_in_slice)
        else:
            assert dlatents_in_slice is None

        if noise is not None:
            del self.top_epi.noise.noise
            gc.collect()

        return x


class InputBlock(nn.Module):
    def __init__(
        self,
        nf,
        dlatent_size,
        const_input_layer,
        gain,
        use_wscale,
        use_noise,
        use_pixel_norm,
        use_instance_norm,
        use_styles,
        activation_layer,
    ):
        super().__init__()
        self.const_input_layer = const_input_layer
        self.nf = nf
        if self.const_input_layer:
            # called 'const' in tf
            self.const = nn.Parameter(th.ones(1, nf, 4, 4))
            self.bias = nn.Parameter(th.ones(nf))
        else:
            self.dense = MyLinear(
                dlatent_size, nf * 16, gain=gain / 4, use_wscale=use_wscale
            )  # tweak gain to match the official implementation of Progressing GAN
        self.epi1 = LayerEpilogue(
            nf, dlatent_size, use_wscale, use_noise, use_pixel_norm, use_instance_norm, use_styles, activation_layer,
        )
        self.conv = MyConv2d(nf, nf, 3, gain=gain, use_wscale=use_wscale)
        self.epi2 = LayerEpilogue(
            nf, dlatent_size, use_wscale, use_noise, use_pixel_norm, use_instance_norm, use_styles, activation_layer,
        )

    def forward(self, dlatents_in_range, noise):
        batch_size = dlatents_in_range.size(0)
        if self.const_input_layer:
            x = self.const.expand(batch_size, -1, -1, -1)
            x = x + self.bias.view(1, -1, 1, 1)
        else:
            x = self.dense(dlatents_in_range[:, 0]).view(batch_size, self.nf, 4, 4)
        x = self.epi1(x, dlatents_in_range[:, 0], noise=noise)
        x = self.conv(x)
        x = self.epi2(x, dlatents_in_range[:, 1], noise=noise)
        return x


class GSynthesisBlock(nn.Module):
    def __init__(
        self,
        in_channels,
        out_channels,
        blur_filter,
        dlatent_size,
        gain,
        use_wscale,
        use_noise,
        use_pixel_norm,
        use_instance_norm,
        use_styles,
        activation_layer,
    ):
        # 2**res x 2**res # res = 3..resolution_log2
        super().__init__()
        if blur_filter:
            blur = BlurLayer(blur_filter)
        else:
            blur = None
        self.conv0_up = MyConv2d(
            in_channels, out_channels, kernel_size=3, gain=gain, use_wscale=use_wscale, intermediate=blur, upscale=True,
        )
        self.epi1 = LayerEpilogue(
            out_channels,
            dlatent_size,
            use_wscale,
            use_noise,
            use_pixel_norm,
            use_instance_norm,
            use_styles,
            activation_layer,
        )
        self.conv1 = MyConv2d(out_channels, out_channels, kernel_size=3, gain=gain, use_wscale=use_wscale)
        self.epi2 = LayerEpilogue(
            out_channels,
            dlatent_size,
            use_wscale,
            use_noise,
            use_pixel_norm,
            use_instance_norm,
            use_styles,
            activation_layer,
        )

    def forward(self, x, dlatents_in_range, noise):
        x = self.conv0_up(x)
        x = self.epi1(x, dlatents_in_range[:, 0], noise=noise)
        x = self.conv1(x)
        x = self.epi2(x, dlatents_in_range[:, 1], noise=noise)
        return x


class G_synthesis(nn.Module):
    def __init__(
        self,
        dlatent_size=512,  # Disentangled latent (W) dimensionality.
        num_channels=3,  # Number of output color channels.
        resolution=1024,  # Output resolution.
        fmap_base=8192,  # Overall multiplier for the number of feature maps.
        fmap_decay=1.0,  # log2 feature map reduction when doubling the resolution.
        fmap_max=512,  # Maximum number of feature maps in any layer.
        use_styles=True,  # Enable style inputs?
        const_input_layer=True,  # First layer is a learned constant?
        use_noise=True,  # Enable noise inputs?
        randomize_noise=False,  # True = randomize noise inputs every time (non-deterministic) or from variables passed,
        nonlinearity="lrelu",  # Activation function: 'relu', 'lrelu'
        use_wscale=True,  # Enable equalized learning rate?
        use_pixel_norm=False,  # Enable pixelwise feature vector normalization?
        use_instance_norm=True,  # Enable instance normalization?
        dtype=th.float32,  # Data type to use for activations and outputs.
        blur_filter=[1, 2, 1],  # Low-pass filter to apply when resampling activations. None = no filtering.
    ):

        super().__init__()

        def nf(stage):
            return min(int(fmap_base / (2.0 ** (stage * fmap_decay))), fmap_max)

        self.dlatent_size = dlatent_size
        resolution_log2 = int(np.log2(resolution))
        assert resolution == 2 ** resolution_log2 and resolution >= 4

        act, gain = {"relu": (th.relu, np.sqrt(2)), "lrelu": (nn.LeakyReLU(negative_slope=0.2), np.sqrt(2)),}[
            nonlinearity
        ]
        blocks = []
        for res in range(2, resolution_log2 + 1):
            channels = nf(res - 1)
            name = "{s}x{s}".format(s=2 ** res)
            if res == 2:
                blocks.append(
                    (
                        name,
                        InputBlock(
                            channels,
                            dlatent_size,
                            const_input_layer,
                            gain,
                            use_wscale,
                            use_noise,
                            use_pixel_norm,
                            use_instance_norm,
                            use_styles,
                            act,
                        ),
                    )
                )

            else:
                blocks.append(
                    (
                        name,
                        GSynthesisBlock(
                            last_channels,
                            channels,
                            blur_filter,
                            dlatent_size,
                            gain,
                            use_wscale,
                            use_noise,
                            use_pixel_norm,
                            use_instance_norm,
                            use_styles,
                            act,
                        ),
                    )
                )
            last_channels = channels
        self.torgb = MyConv2d(channels, num_channels, 1, gain=1, use_wscale=use_wscale)
        self.blocks = nn.ModuleDict(OrderedDict(blocks))

    def forward(self, dlatents_in, noise):
        # Input: Disentangled latents (W) [minibatch, num_layers, dlatent_size].
        # lod_in = tf.cast(tf.get_variable('lod', initializer=np.float32(0), trainable=False), dtype)
        for i, m in enumerate(self.blocks.values()):
            if i == 0:
                x = m(dlatents_in[:, 2 * i : 2 * i + 2], noise=noise)
            else:
                x = m(x, dlatents_in[:, 2 * i : 2 * i + 2], noise=noise)
        rgb = self.torgb(x)
        return rgb


class G_style(nn.Sequential):
    def __init__(self, output_size=1920, checkpoint=None):

        # TODO FIX THIS MONSTROSITY
        super().__init__()
        self.g_mapping = G_mapping()
        try:
            self.g_synthesis = G_synthesis(resolution=1024)
            if checkpoint is not None:
                self.load_state_dict(th.load(checkpoint), strict=False)
            network_resolution = 1024
        except:
            print("Trying 512px generator resolution...")
            try:
                self.g_synthesis = G_synthesis(resolution=512)
                if checkpoint is not None:
                    self.load_state_dict(th.load(checkpoint), strict=False)
                network_resolution = 512
            except:
                print("Trying 256px generator resolution...")
                try:
                    self.g_synthesis = G_synthesis(resolution=256)
                    if checkpoint is not None:
                        self.load_state_dict(th.load(checkpoint), strict=False)
                    network_resolution = 256
                except:
                    print("Trying 128px generator resolution...")
                    try:
                        self.g_synthesis = G_synthesis(resolution=128)
                        if checkpoint is not None:
                            self.load_state_dict(th.load(checkpoint), strict=False)
                        network_resolution = 128
                    except:
                        print("ERROR: Network too small or state_dict mismatch")
                        exit()

        const = getattr(self.g_synthesis.blocks, "4x4").const
        if network_resolution != 1024:
            means = th.zeros(size=(1, 512, int(4 * 1024 / network_resolution), int(4 * 1024 / network_resolution)))
            const = th.normal(mean=means, std=th.ones_like(means) * const.std(),)

        _, _, ch, cw = const.shape
        if output_size == 1920:
            layer0 = th.cat(
                [
                    const[:, :, :, [0]],
                    const[:, :, :, [0]],
                    # const[:, :, :, : cw // 2 + 1][:, :, :, list(range(cw // 2, 0, -1))],
                    const,
                    # const[:, :, :, cw // 2 :],
                    const[:, :, :, [-1]],
                    const[:, :, :, [-1]],
                ],
                axis=3,
            )
        elif output_size == 512:
            layer0 = const[:, :, ch // 4 : 3 * ch // 4, cw // 4 : 3 * cw // 4]
        else:
            layer0 = const
        getattr(self.g_synthesis.blocks, "4x4").const = th.nn.Parameter(layer0 + th.normal(0, const.std() / 2.0))
        _, _, height, width = getattr(self.g_synthesis.blocks, "4x4").const.shape

        for i in range(len(list(self.g_synthesis.blocks.named_parameters())) // 10):
            self.register_buffer(f"noise_{i}", th.randn(1, 1, height * 2 ** i, width * 2 ** i))

        self.truncation_latent = self.mean_latent(2 ** 14)

    def mean_latent(self, n_latent):
        latent_in = th.randn(n_latent, 512)
        latent = self.g_mapping(latent_in).mean(0, keepdim=True)
        return latent

    def forward(
        self,
        styles,
        noise=None,
        truncation=1,
        map_latents=False,
        randomize_noise=False,
        input_is_latent=True,
        transform_dict_list=None,
    ):
        if map_latents:
            return self.g_mapping(styles)

        if noise is None:
            noise = [None] * (len(list(self.g_synthesis.blocks.named_parameters())) // 10)
        for ns, noise_scale in enumerate(noise):
            if noise_scale is None:
                try:
                    noise[ns] = getattr(self, f"noise_{ns}")
                except:
                    pass

        if truncation != 1:
            interp = th.lerp(self.truncation_latent.to(styles.device), styles, truncation)
            do_trunc = (th.arange(styles.size(1)) < 8).view(1, -1, 1).to(styles.device)
            styles = th.where(do_trunc, interp, styles)

        # Input: Disentangled latents (W) [minibatch, num_layers, dlatent_size].
        # print(styles.shape, len(noise), len(self.g_synthesis.blocks.values()))
        for i, block in enumerate(self.g_synthesis.blocks.values()):
            if i == 0:
                x = block(styles[:, 2 * i : 2 * i + 2], noise=noise[i])
            else:
                x = block(x, styles[:, 2 * i : 2 * i + 2], noise=noise[i])
        img = self.g_synthesis.torgb(x)

        return img, None


================================================
FILE: models/stylegan2.py
================================================
import math
import os
import random
import sys

import torch as th
from torch import nn
from torch.nn import functional as F

sys.path.append(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))

from op import FusedLeakyReLU, fused_leaky_relu, upfirdn2d


class PixelNorm(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, inputs):
        return inputs * th.rsqrt(th.mean(inputs ** 2, dim=1, keepdim=True) + 1e-8)


def make_kernel(k):
    k = th.tensor(k, dtype=th.float32)

    if k.ndim == 1:
        k = k[None, :] * k[:, None]

    k /= k.sum()

    return k


class Upsample(nn.Module):
    def __init__(self, kernel, factor=2):
        super().__init__()

        self.factor = factor
        kernel = make_kernel(kernel) * (factor ** 2)
        self.register_buffer("kernel", kernel)

        p = kernel.shape[0] - factor

        pad0 = (p + 1) // 2 + factor - 1
        pad1 = p // 2

        self.pad = (pad0, pad1)

    def forward(self, inputs):
        out = upfirdn2d(inputs, self.kernel, up=self.factor, down=1, pad=self.pad)

        return out


class Downsample(nn.Module):
    def __init__(self, kernel, factor=2):
        super().__init__()

        self.factor = factor
        kernel = make_kernel(kernel)
        self.register_buffer("kernel", kernel)

        p = kernel.shape[0] - factor

        pad0 = (p + 1) // 2
        pad1 = p // 2

        self.pad = (pad0, pad1)

    def forward(self, inputs):
        out = upfirdn2d(inputs, self.kernel, up=1, down=self.factor, pad=self.pad)

        return out


class Blur(nn.Module):
    def __init__(self, kernel, pad, upsample_factor=1):
        super().__init__()

        kernel = make_kernel(kernel)

        if upsample_factor > 1:
            kernel = kernel * (upsample_factor ** 2)

        self.register_buffer("kernel", kernel)

        self.pad = pad

    def forward(self, inputs):
        out = upfirdn2d(inputs, self.kernel, pad=self.pad)

        return out


class EqualConv2d(nn.Module):
    def __init__(self, in_channel, out_channel, kernel_size, stride=1, padding=0, bias=True):
        super().__init__()

        self.weight = nn.Parameter(th.randn(out_channel, in_channel, kernel_size, kernel_size))
        self.scale = 1 / math.sqrt(in_channel * kernel_size ** 2)

        self.stride = stride
        self.padding = padding

        if bias:
            self.bias = nn.Parameter(th.zeros(out_channel))

        else:
            self.bias = None

    def forward(self, inputs):
        out = F.conv2d(inputs, self.weight * self.scale, bias=self.bias, stride=self.stride, padding=self.padding,)

        return out

    def __repr__(self):
        return (
            f"{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]},"
            f" {self.weight.shape[2]}, stride={self.stride}, padding={self.padding})"
        )


class EqualLinear(nn.Module):
    def __init__(self, in_dim, out_dim, bias=True, bias_init=0, lr_mul=1, activation=None):
        super().__init__()

        self.weight = nn.Parameter(th.randn(out_dim, in_dim).div_(lr_mul))

        if bias:
            self.bias = nn.Parameter(th.zeros(out_dim).fill_(bias_init))

        else:
            self.bias = None

        self.activation = activation

        self.scale = (1 / math.sqrt(in_dim)) * lr_mul
        self.lr_mul = lr_mul

    def forward(self, inputs):
        if self.activation:
            out = F.linear(inputs, self.weight * self.scale)
            out = fused_leaky_relu(out, self.bias * self.lr_mul)
        else:
            out = F.linear(inputs, self.weight * self.scale, bias=self.bias * self.lr_mul)
        return out

    def __repr__(self):
        return f"{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]})"


class ScaledLeakyReLU(nn.Module):
    def __init__(self, negative_slope=0.2):
        super().__init__()

        self.negative_slope = negative_slope

    def forward(self, inputs):
        out = F.leaky_relu(inputs, negative_slope=self.negative_slope)

        return out * math.sqrt(2)


class ModulatedConv2d(nn.Module):
    def __init__(
        self,
        in_channel,
        out_channel,
        kernel_size,
        style_dim,
        demodulate=True,
        upsample=False,
        downsample=False,
        blur_kernel=[1, 3, 3, 1],
    ):
        super().__init__()

        self.eps = 1e-8
        self.kernel_size = kernel_size
        self.in_channel = in_channel
        self.out_channel = out_channel
        self.upsample = upsample
        self.downsample = downsample

        if upsample:
            factor = 2
            p = (len(blur_kernel) - factor) - (kernel_size - 1)
            pad0 = (p + 1) // 2 + factor - 1
            pad1 = p // 2 + 1

            self.blur = Blur(blur_kernel, pad=(pad0, pad1), upsample_factor=factor)

        if downsample:
            factor = 2
            p = (len(blur_kernel) - factor) + (kernel_size - 1)
            pad0 = (p + 1) // 2
            pad1 = p // 2

            self.blur = Blur(blur_kernel, pad=(pad0, pad1))

        fan_in = in_channel * kernel_size ** 2
        self.scale = 1 / math.sqrt(fan_in)
        self.padding = kernel_size // 2

        self.weight = nn.Parameter(th.randn(1, out_channel, in_channel, kernel_size, kernel_size))

        self.modulation = EqualLinear(style_dim, in_channel, bias_init=1)

        self.demodulate = demodulate

    def __repr__(self):
        return (
            f"{self.__class__.__name__}({self.in_channel}, {self.out_channel}, {self.kernel_size}, "
            f"upsample={self.upsample}, downsample={self.downsample})"
        )

    def forward(self, inputs, style):
        batch, in_channel, height, width = inputs.shape

        style = self.modulation(style).view(batch, 1, in_channel, 1, 1)
        weight = self.scale * self.weight * style

        if self.demodulate:
            demod = th.rsqrt(weight.pow(2).sum([2, 3, 4]) + 1e-8)
            weight = weight * demod.view(batch, self.out_channel, 1, 1, 1)

        weight = weight.view(batch * self.out_channel, in_channel, self.kernel_size, self.kernel_size)

        if self.upsample:
            inputs = inputs.view(1, batch * in_channel, height, width)
            weight = weight.view(batch, self.out_channel, in_channel, self.kernel_size, self.kernel_size)
            weight = weight.transpose(1, 2).reshape(
                batch * in_channel, self.out_channel, self.kernel_size, self.kernel_size
            )
            out = F.conv_transpose2d(inputs, weight, padding=0, stride=2, groups=batch)
            _, _, height, width = out.shape
            out = out.view(batch, self.out_channel, height, width)
            out = self.blur(out)

        elif self.downsample:
            inputs = self.blur(inputs)
            _, _, height, width = inputs.shape
            inputs = inputs.view(1, batch * in_channel, height, width)
            out = F.conv2d(inputs, weight, padding=0, stride=2, groups=batch)
            _, _, height, width = out.shape
            out = out.view(batch, self.out_channel, height, width)

        else:
            inputs = inputs.view(1, batch * in_channel, height, width)
            out = F.conv2d(inputs, weight, padding=self.padding, groups=batch)
            _, _, height, width = out.shape
            out = out.view(batch, self.out_channel, height, width)

        return out


class NoiseInjection(nn.Module):
    def __init__(self):
        super().__init__()
        self.weight = nn.Parameter(th.zeros(1))

    def forward(self, image, noise=None):
        if noise is None:
            batch, _, height, width = image.shape
            noise = image.new_empty(batch, 1, height, width).normal_()
        return image + self.weight * noise.to(image.device)


class ConstantInput(nn.Module):
    def __init__(self, channel, size=4):
        super().__init__()

        self.input = nn.Parameter(th.randn(1, channel, size, size))

    def forward(self, inputs):
        batch = inputs.shape[0]
        out = self.input.repeat(batch, 1, 1, 1)
        return out


class LatentInput(nn.Module):
    def __init__(self, latent_dim, channel, size=4):
        super().__init__()
        self.channel = channel
        self.size = size
        self.linear = EqualLinear(latent_dim, channel * size * size, activation="fused_lrelu")
        self.activate = FusedLeakyReLU(channel * size * size)
        self.input = nn.Parameter(th.randn(1))

    def forward(self, inputs):
        batch = inputs.shape[0]
        out = self.linear(inputs[:, 0])
        out = self.activate(out)
        return out.reshape((batch, self.channel, self.size, self.size))


class ManipulationLayer(th.nn.Module):
    def __init__(self, layer):
        super().__init__()
        self.layer = layer

    def forward(self, input, tranforms_dict_list):
        out = input
        for transform_dict in tranforms_dict_list:
            if transform_dict["layer"] == self.layer:
                out = transform_dict["transform"].to(out.device)(out)
        return out


class StyledConv(nn.Module):
    def __init__(
        self,
        in_channel,
        out_channel,
        kernel_size,
        style_dim,
        upsample=False,
        blur_kernel=[1, 3, 3, 1],
        demodulate=True,
        layerID=-1,
    ):
        super().__init__()

        self.conv = ModulatedConv2d(
            in_channel,
            out_channel,
            kernel_size,
            style_dim,
            upsample=upsample,
            blur_kernel=blur_kernel,
            demodulate=demodulate,
        )

        self.noise = NoiseInjection()
        self.activate = FusedLeakyReLU(out_channel)
        self.manipulation = ManipulationLayer(layerID)

    def forward(self, inputs, style, noise=None, transform_dict_list=[]):
        out = self.conv(inputs, style)
        out = self.noise(out, noise=noise)
        out = self.activate(out)
        out = self.manipulation(out, transform_dict_list)
        return out


class ToRGB(nn.Module):
    def __init__(self, in_channel, style_dim, upsample=True, blur_kernel=[1, 3, 3, 1]):
        super().__init__()

        if upsample:
            self.upsample = Upsample(blur_kernel)

        self.conv = ModulatedConv2d(in_channel, 3, 1, style_dim, demodulate=False)
        self.bias = nn.Parameter(th.zeros(1, 3, 1, 1))

    def forward(self, inputs, style, skip=None):
        out = self.conv(inputs, style)
        out = out + self.bias

        if skip is not None:
            skip = self.upsample(skip)

            out = out + skip

        return out


class Generator(nn.Module):
    def __init__(
        self,
        size,
        style_dim,
        n_mlp,
        channel_multiplier=2,
        blur_kernel=[1, 3, 3, 1],
        lr_mlp=0.01,
        constant_input=False,
        checkpoint=None,
        output_size=None,
        min_rgb_size=4,
        base_res_factor=1,
    ):
        super().__init__()

        self.size = size
        self.style_dim = style_dim

        layers = [PixelNorm()]

        for i in range(n_mlp):
            layers.append(EqualLinear(style_dim, style_dim, lr_mul=lr_mlp, activation="fused_lrelu"))

        self.style = nn.Sequential(*layers)

        self.channels = {
            4: 512,
            8: 512,
            16: 512,
            32: 512,
            64: 256 * channel_multiplier,
            128: 128 * channel_multiplier,
            256: 64 * channel_multiplier,
            512: 32 * channel_multiplier,
            1024: 16 * channel_multiplier,
        }

        self.log_size = int(math.log(size, 2))
        self.num_layers = (self.log_size - 2) * 2 + 1
        self.n_latent = self.log_size * 2 - 2
        self.min_rgb_size = min_rgb_size

        if constant_input:
            self.input = ConstantInput(self.channels[4])
        else:
            self.input = LatentInput(style_dim, self.channels[4])

        self.const_manipulation = ManipulationLayer(0)

        layerID = 1
        self.conv1 = StyledConv(
            self.channels[4], self.channels[4], 3, style_dim, blur_kernel=blur_kernel, layerID=layerID
        )
        self.to_rgb1 = ToRGB(self.channels[4], style_dim, upsample=False)

        self.convs = nn.ModuleList()
        self.upsamples = nn.ModuleList()
        self.to_rgbs = nn.ModuleList()
        self.noises = nn.Module()

        in_channel = self.channels[4]

        for layer_idx in range(self.num_layers):
            res = (layer_idx + 5) // 2
            shape = [1, 1, 2 ** res, 2 ** res]
            self.noises.register_buffer(f"noise_{layer_idx}", th.randn(*shape))

        for i in range(3, self.log_size + 1):
            out_channel = self.channels[2 ** i]

            layerID += 1
            self.convs.append(
                StyledConv(
                    in_channel, out_channel, 3, style_dim, upsample=True, blur_kernel=blur_kernel, layerID=layerID
                )
            )

            layerID += 1
            self.convs.append(
                StyledConv(out_channel, out_channel, 3, style_dim, blur_kernel=blur_kernel, layerID=layerID)
            )

            self.to_rgbs.append(ToRGB(out_channel, style_dim))

            in_channel = out_channel

        self.truncation_latent = None

        if checkpoint is not None:
            self.load_state_dict(th.load(checkpoint)["g_ema"])

        if size != output_size or base_res_factor != 1:
            for layer_idx in range(self.num_layers):
                res = (layer_idx + 5) // 2
                shape = [
                    1,
                    1,
                    int(base_res_factor * 2 ** res * (2 if output_size == 1080 else 1)),
                    int(base_res_factor * 2 ** res * (2 if output_size == 1920 else 1)),
                ]
                setattr(self.noises, f"noise_{layer_idx}", th.randn(*shape))

    def make_noise(self):
        device = self.input.input.device

        noises = [th.randn(1, 1, 2 ** 2, 2 ** 2, device=device)]

        for i in range(3, self.log_size + 1):
            for _ in range(2):
                noises.append(th.randn(1, 1, 2 ** i, 2 ** i, device=device))

        return noises

    def mean_latent(self, n_latent):
        latent_in = th.randn(n_latent, self.style_dim, device=self.input.input.device)
        latent = self.style(latent_in).mean(0, keepdim=True)

        return latent

    def get_latent(self, inputs):
        return self.style(inputs)

    def forward(
        self,
        styles,
        return_latents=False,
        return_activation_maps=False,
        inject_index=None,
        truncation=1.0,
        truncation_latent=None,
        input_is_latent=False,
        noise=None,
        randomize_noise=True,
        transform_dict_list=[],
        map_latents=False,
    ):
        if map_latents:
            latent = th.cat([self.style(s[None, None, :]) for s in styles], axis=0)
            latent = latent.repeat(1, self.n_latent, 1)
            return latent

        if not input_is_latent:
            styles = [self.style(s) for s in styles]

            if len(styles) < 2:
                inject_index = self.n_latent
                if styles[0].ndim < 3:
                    latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1)
                else:
                    latent = styles[0]
            else:
                if inject_index is None:
                    inject_index = random.randint(1, self.n_latent - 1)
                latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1)
                latent2 = styles[1].unsqueeze(1).repeat(1, self.n_latent - inject_index, 1)
                latent = th.cat([latent, latent2], 1)
        else:
            latent = styles
            if latent.dim() == 2:
                latent = latent[:, None, :].repeat(1, self.n_latent, 1)

        if noise is None:
            noise = [None] * self.num_layers
        for ns, noise_scale in enumerate(noise):
            if not randomize_noise and noise_scale is None:
                noise[ns] = getattr(self.noises, f"noise_{ns}")

        if isinstance(truncation, float):
            truncation = th.cuda.FloatTensor([truncation])
        if self.truncation_latent is None:
            self.truncation_latent = truncation_latent if truncation_latent is not None else self.mean_latent(2 ** 14)
        latent = self.truncation_latent[None, ...] + truncation.to(latent.device)[:, None, None] * (
            latent - self.truncation_latent[None, ...]
        )

        activation_map_list = []

        out = self.input(latent)
        out = self.const_manipulation(out, transform_dict_list)
        out = self.conv1(out, latent[:, 0], noise=noise[0], transform_dict_list=transform_dict_list)
        activation_map_list.append(out)

        current_size = 4
        if self.min_rgb_size <= current_size:
            image = self.to_rgb1(out, latent[:, 1])
        else:
            image = None

        i = 1
        for conv1, conv2, noise1, noise2, to_rgb in zip(
            self.convs[::2], self.convs[1::2], noise[1::2], noise[2::2], self.to_rgbs
        ):
            out = conv1(out, latent[:, i], noise=noise1, transform_dict_list=transform_dict_list)
            current_size *= 2
            activation_map_list.append(out)
            out = conv2(out, latent[:, i + 1], noise=noise2, transform_dict_list=transform_dict_list)
            activation_map_list.append(out)
            if self.min_rgb_size <= current_size:
                image = to_rgb(out, latent[:, i + 2], image)
            i += 2

        if return_activation_maps:
            return image, activation_map_list
        elif return_latents:
            return image, latent
        else:
            return image, None


class ConvLayer(nn.Sequential):
    def __init__(
        self,
        in_channel,
        out_channel,
        kernel_size,
        downsample=False,
        blur_kernel=[1, 3, 3, 1],
        bias=True,
        activate=True,
    ):
        layers = []

        if downsample:
            factor = 2
            p = (len(blur_kernel) - factor) + (kernel_size - 1)
            pad0 = (p + 1) // 2
            pad1 = p // 2

            layers.append(Blur(blur_kernel, pad=(pad0, pad1)))

            stride = 2
            self.padding = 0

        else:
            stride = 1
            self.padding = kernel_size // 2

        layers.append(
            EqualConv2d(
                in_channel, out_channel, kernel_size, padding=self.padding, stride=stride, bias=bias and not activate,
            )
        )

        if activate:
            if bias:
                layers.append(FusedLeakyReLU(out_channel))

            else:
                layers.append(ScaledLeakyReLU(0.2))

        super().__init__(*layers)


class ResBlock(nn.Module):
    def __init__(self, in_channel, out_channel, blur_kernel=[1, 3, 3, 1], use_skip=True):
        super().__init__()

        self.conv1 = ConvLayer(in_channel, in_channel, 3)
        self.conv2 = ConvLayer(in_channel, out_channel, 3, downsample=True)

        if use_skip:
            self.skip = ConvLayer(in_channel, out_channel, 1, downsample=True, activate=False, bias=False)
        else:
            self.skip = None

    def forward(self, inputs):
        out = self.conv1(inputs)
        out = self.conv2(out)

        if self.skip is not None:
            skip = self.skip(inputs)
            out = (out + skip) / math.sqrt(2)

        return out


class Discriminator(nn.Module):
    def __init__(self, size, channel_multiplier=2, blur_kernel=[1, 3, 3, 1], use_skip=True):
        super().__init__()

        channels = {
            4: 512,
            8: 512,
            16: 512,
            32: 512,
            64: 256 * channel_multiplier,
            128: 128 * channel_multiplier,
            256: 64 * channel_multiplier,
            512: 32 * channel_multiplier,
            1024: 16 * channel_multiplier,
        }

        convs = [ConvLayer(3, channels[size], 1)]

        log_size = int(math.log(size, 2))

        in_channel = channels[size]

        for i in range(log_size, 2, -1):
            out_channel = channels[2 ** (i - 1)]

            convs.append(ResBlock(in_channel, out_channel, blur_kernel, use_skip=use_skip))

            in_channel = out_channel

        self.convs = nn.Sequential(*convs)

        self.stddev_group = 4
        self.stddev_feat = 1

        self.final_conv = ConvLayer(in_channel + 1, channels[4], 3)
        self.final_linear = nn.Sequential(
            EqualLinear(channels[4] * 4 * 4, channels[4], activation="fused_lrelu"), EqualLinear(channels[4], 1),
        )

    def forward(self, inputs):
        out = self.convs(inputs)

        batch, channel, height, width = out.shape

        try:
            group = min(batch, self.stddev_group)
            stddev = out.view(group, -1, self.stddev_feat, channel // self.stddev_feat, height, width)
            stddev = th.sqrt(stddev.var(0, unbiased=False) + 1e-8)
            stddev = stddev.mean([2, 3, 4], keepdims=True).squeeze(2)
            stddev = stddev.repeat(group, 1, height, width)
            out = th.cat([out, stddev], 1)
        except RuntimeError:
            group = batch
            stddev = out.view(group, -1, self.stddev_feat, channel // self.stddev_feat, height, width)
            stddev = th.sqrt(stddev.var(0, unbiased=False) + 1e-8)
            stddev = stddev.mean([2, 3, 4], keepdims=True).squeeze(2)
            stddev = stddev.repeat(group, 1, height, width)
            out = th.cat([out, stddev], 1)

        out = self.final_conv(out)

        out = out.view(batch, -1)
        out = self.final_linear(out)

        return out


================================================
FILE: op/__init__.py
================================================
from .fused_act import FusedLeakyReLU, fused_leaky_relu
from .upfirdn2d import upfirdn2d


================================================
FILE: op/fused_act.py
================================================
import os

import torch
from torch import nn
from torch.nn import functional as F
from torch.autograd import Function
from torch.utils.cpp_extension import load


module_path = os.path.dirname(__file__)
fused = load(
    "fused",
    sources=[
        os.path.join(module_path, "fused_bias_act.cpp"),
        os.path.join(module_path, "fused_bias_act_kernel.cu"),
    ],
)


class FusedLeakyReLUFunctionBackward(Function):
    @staticmethod
    def forward(ctx, grad_output, out, negative_slope, scale):
        ctx.save_for_backward(out)
        ctx.negative_slope = negative_slope
        ctx.scale = scale

        empty = grad_output.new_empty(0)

        grad_input = fused.fused_bias_act(
            grad_output, empty, out, 3, 1, negative_slope, scale
        )

        dim = [0]

        if grad_input.ndim > 2:
            dim += list(range(2, grad_input.ndim))

        grad_bias = grad_input.sum(dim).detach()

        return grad_input, grad_bias

    @staticmethod
    def backward(ctx, gradgrad_input, gradgrad_bias):
        out, = ctx.saved_tensors
        gradgrad_out = fused.fused_bias_act(
            gradgrad_input, gradgrad_bias, out, 3, 1, ctx.negative_slope, ctx.scale
        )

        return gradgrad_out, None, None, None


class FusedLeakyReLUFunction(Function):
    @staticmethod
    def forward(ctx, input, bias, negative_slope, scale):
        empty = input.new_empty(0)
        out = fused.fused_bias_act(input, bias, empty, 3, 0, negative_slope, scale)
        ctx.save_for_backward(out)
        ctx.negative_slope = negative_slope
        ctx.scale = scale

        return out

    @staticmethod
    def backward(ctx, grad_output):
        out, = ctx.saved_tensors

        grad_input, grad_bias = FusedLeakyReLUFunctionBackward.apply(
            grad_output, out, ctx.negative_slope, ctx.scale
        )

        return grad_input, grad_bias, None, None


class FusedLeakyReLU(nn.Module):
    def __init__(self, channel, negative_slope=0.2, scale=2 ** 0.5):
        super().__init__()

        self.bias = nn.Parameter(torch.zeros(channel))
        self.negative_slope = negative_slope
        self.scale = scale

    def forward(self, input):
        return fused_leaky_relu(input, self.bias, self.negative_slope, self.scale)


def fused_leaky_relu(input, bias, negative_slope=0.2, scale=2 ** 0.5):
    if input.device.type == "cpu":
        rest_dim = [1] * (input.ndim - bias.ndim - 1)
        return (
            F.leaky_relu(
                input + bias.view(1, bias.shape[0], *rest_dim), negative_slope=0.2
            )
            * scale
        )

    else:
        return FusedLeakyReLUFunction.apply(input, bias, negative_slope, scale)


================================================
FILE: op/fused_bias_act.cpp
================================================
#include <torch/extension.h>


torch::Tensor fused_bias_act_op(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer,
    int act, int grad, float alpha, float scale);

#define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)

torch::Tensor fused_bias_act(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer,
    int act, int grad, float alpha, float scale) {
    CHECK_CUDA(input);
    CHECK_CUDA(bias);

    return fused_bias_act_op(input, bias, refer, act, grad, alpha, scale);
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("fused_bias_act", &fused_bias_act, "fused bias act (CUDA)");
}

================================================
FILE: op/fused_bias_act_kernel.cu
================================================
// 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

#include <torch/types.h>

#include <ATen/ATen.h>
#include <ATen/AccumulateType.h>
#include <ATen/cuda/CUDAContext.h>
#include <ATen/cuda/CUDAApplyUtils.cuh>

#include <cuda.h>
#include <cuda_runtime.h>


template <typename scalar_t>
static __global__ void fused_bias_act_kernel(scalar_t* out, const scalar_t* p_x, const scalar_t* p_b, const scalar_t* p_ref,
    int act, int grad, scalar_t alpha, scalar_t scale, int loop_x, int size_x, int step_b, int size_b, int use_bias, int use_ref) {
    int xi = blockIdx.x * loop_x * blockDim.x + threadIdx.x;

    scalar_t zero = 0.0;

    for (int loop_idx = 0; loop_idx < loop_x && xi < size_x; loop_idx++, xi += blockDim.x) {
        scalar_t x = p_x[xi];

        if (use_bias) {
            x += p_b[(xi / step_b) % size_b];
        }

        scalar_t ref = use_ref ? p_ref[xi] : zero;

        scalar_t y;

        switch (act * 10 + grad) {
            default:
            case 10: y = x; break;
            case 11: y = x; break;
            case 12: y = 0.0; break;

            case 30: y = (x > 0.0) ? x : x * alpha; break;
            case 31: y = (ref > 0.0) ? x : x * alpha; break;
            case 32: y = 0.0; break;
        }

        out[xi] = y * scale;
    }
}


torch::Tensor fused_bias_act_op(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer,
    int act, int grad, float alpha, float scale) {
    int curDevice = -1;
    cudaGetDevice(&curDevice);
    cudaStream_t stream = at::cuda::getCurrentCUDAStream(curDevice);

    auto x = input.contiguous();
    auto b = bias.contiguous();
    auto ref = refer.contiguous();

    int use_bias = b.numel() ? 1 : 0;
    int use_ref = ref.numel() ? 1 : 0;

    int size_x = x.numel();
    int size_b = b.numel();
    int step_b = 1;

    for (int i = 1 + 1; i < x.dim(); i++) {
        step_b *= x.size(i);
    }

    int loop_x = 4;
    int block_size = 4 * 32;
    int grid_size = (size_x - 1) / (loop_x * block_size) + 1;

    auto y = torch::empty_like(x);

    AT_DISPATCH_FLOATING_TYPES_AND_HALF(x.scalar_type(), "fused_bias_act_kernel", [&] {
        fused_bias_act_kernel<scalar_t><<<grid_size, block_size, 0, stream>>>(
            y.data_ptr<scalar_t>(),
            x.data_ptr<scalar_t>(),
            b.data_ptr<scalar_t>(),
            ref.data_ptr<scalar_t>(),
            act,
            grad,
            alpha,
            scale,
            loop_x,
            size_x,
            step_b,
            size_b,
            use_bias,
            use_ref
        );
    });

    return y;
}

================================================
FILE: op/upfirdn2d.cpp
================================================
#include <torch/extension.h>


torch::Tensor upfirdn2d_op(const torch::Tensor& input, const torch::Tensor& kernel,
                            int up_x, int up_y, int down_x, int down_y,
                            int pad_x0, int pad_x1, int pad_y0, int pad_y1);

#define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor")
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)

torch::Tensor upfirdn2d(const torch::Tensor& input, const torch::Tensor& kernel,
                        int up_x, int up_y, int down_x, int down_y,
                        int pad_x0, int pad_x1, int pad_y0, int pad_y1) {
    CHECK_CUDA(input);
    CHECK_CUDA(kernel);

    return upfirdn2d_op(input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1);
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("upfirdn2d", &upfirdn2d, "upfirdn2d (CUDA)");
}

================================================
FILE: op/upfirdn2d.py
================================================
import os

import torch
from torch.nn import functional as F
from torch.autograd import Function
from torch.utils.cpp_extension import load


module_path = os.path.dirname(__file__)
upfirdn2d_op = load(
    "upfirdn2d",
    sources=[
        os.path.join(module_path, "upfirdn2d.cpp"),
        os.path.join(module_path, "upfirdn2d_kernel.cu"),
    ],
)


class UpFirDn2dBackward(Function):
    @staticmethod
    def forward(
        ctx, grad_output, kernel, grad_kernel, up, down, pad, g_pad, in_size, out_size
    ):

        up_x, up_y = up
        down_x, down_y = down
        g_pad_x0, g_pad_x1, g_pad_y0, g_pad_y1 = g_pad

        grad_output = grad_output.reshape(-1, out_size[0], out_size[1], 1)

        grad_input = upfirdn2d_op.upfirdn2d(
            grad_output,
            grad_kernel,
            down_x,
            down_y,
            up_x,
            up_y,
            g_pad_x0,
            g_pad_x1,
            g_pad_y0,
            g_pad_y1,
        )
        grad_input = grad_input.view(in_size[0], in_size[1], in_size[2], in_size[3])

        ctx.save_for_backward(kernel)

        pad_x0, pad_x1, pad_y0, pad_y1 = pad

        ctx.up_x = up_x
        ctx.up_y = up_y
        ctx.down_x = down_x
        ctx.down_y = down_y
        ctx.pad_x0 = pad_x0
        ctx.pad_x1 = pad_x1
        ctx.pad_y0 = pad_y0
        ctx.pad_y1 = pad_y1
        ctx.in_size = in_size
        ctx.out_size = out_size

        return grad_input

    @staticmethod
    def backward(ctx, gradgrad_input):
        kernel, = ctx.saved_tensors

        gradgrad_input = gradgrad_input.reshape(-1, ctx.in_size[2], ctx.in_size[3], 1)

        gradgrad_out = upfirdn2d_op.upfirdn2d(
            gradgrad_input,
            kernel,
            ctx.up_x,
            ctx.up_y,
            ctx.down_x,
            ctx.down_y,
            ctx.pad_x0,
            ctx.pad_x1,
            ctx.pad_y0,
            ctx.pad_y1,
        )
        # gradgrad_out = gradgrad_out.view(ctx.in_size[0], ctx.out_size[0], ctx.out_size[1], ctx.in_size[3])
        gradgrad_out = gradgrad_out.view(
            ctx.in_size[0], ctx.in_size[1], ctx.out_size[0], ctx.out_size[1]
        )

        return gradgrad_out, None, None, None, None, None, None, None, None


class UpFirDn2d(Function):
    @staticmethod
    def forward(ctx, input, kernel, up, down, pad):
        up_x, up_y = up
        down_x, down_y = down
        pad_x0, pad_x1, pad_y0, pad_y1 = pad

        kernel_h, kernel_w = kernel.shape
        batch, channel, in_h, in_w = input.shape
        ctx.in_size = input.shape

        input = input.reshape(-1, in_h, in_w, 1)

        ctx.save_for_backward(kernel, torch.flip(kernel, [0, 1]))

        out_h = (in_h * up_y + pad_y0 + pad_y1 - kernel_h) // down_y + 1
        out_w = (in_w * up_x + pad_x0 + pad_x1 - kernel_w) // down_x + 1
        ctx.out_size = (out_h, out_w)

        ctx.up = (up_x, up_y)
        ctx.down = (down_x, down_y)
        ctx.pad = (pad_x0, pad_x1, pad_y0, pad_y1)

        g_pad_x0 = kernel_w - pad_x0 - 1
        g_pad_y0 = kernel_h - pad_y0 - 1
        g_pad_x1 = in_w * up_x - out_w * down_x + pad_x0 - up_x + 1
        g_pad_y1 = in_h * up_y - out_h * down_y + pad_y0 - up_y + 1

        ctx.g_pad = (g_pad_x0, g_pad_x1, g_pad_y0, g_pad_y1)

        out = upfirdn2d_op.upfirdn2d(
            input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1
        )
        # out = out.view(major, out_h, out_w, minor)
        out = out.view(-1, channel, out_h, out_w)

        return out

    @staticmethod
    def backward(ctx, grad_output):
        kernel, grad_kernel = ctx.saved_tensors

        grad_input = UpFirDn2dBackward.apply(
            grad_output,
            kernel,
            grad_kernel,
            ctx.up,
            ctx.down,
            ctx.pad,
            ctx.g_pad,
            ctx.in_size,
            ctx.out_size,
        )

        return grad_input, None, None, None, None


def upfirdn2d(input, kernel, up=1, down=1, pad=(0, 0)):
    if input.device.type == "cpu":
        out = upfirdn2d_native(
            input, kernel, up, up, down, down, pad[0], pad[1], pad[0], pad[1]
        )

    else:
        out = UpFirDn2d.apply(
            input, kernel, (up, up), (down, down), (pad[0], pad[1], pad[0], pad[1])
        )

    return out


def upfirdn2d_native(
    input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1
):
    _, channel, in_h, in_w = input.shape
    input = input.reshape(-1, in_h, in_w, 1)

    _, in_h, in_w, minor = input.shape
    kernel_h, kernel_w = kernel.shape

    out = input.view(-1, in_h, 1, in_w, 1, minor)
    out = F.pad(out, [0, 0, 0, up_x - 1, 0, 0, 0, up_y - 1])
    out = out.view(-1, in_h * up_y, in_w * up_x, minor)

    out = F.pad(
        out, [0, 0, max(pad_x0, 0), max(pad_x1, 0), max(pad_y0, 0), max(pad_y1, 0)]
    )
    out = out[
        :,
        max(-pad_y0, 0) : out.shape[1] - max(-pad_y1, 0),
        max(-pad_x0, 0) : out.shape[2] - max(-pad_x1, 0),
        :,
    ]

    out = out.permute(0, 3, 1, 2)
    out = out.reshape(
        [-1, 1, in_h * up_y + pad_y0 + pad_y1, in_w * up_x + pad_x0 + pad_x1]
    )
    w = torch.flip(kernel, [0, 1]).view(1, 1, kernel_h, kernel_w)
    out = F.conv2d(out, w)
    out = out.reshape(
        -1,
        minor,
        in_h * up_y + pad_y0 + pad_y1 - kernel_h + 1,
        in_w * up_x + pad_x0 + pad_x1 - kernel_w + 1,
    )
    out = out.permute(0, 2, 3, 1)
    out = out[:, ::down_y, ::down_x, :]

    out_h = (in_h * up_y + pad_y0 + pad_y1 - kernel_h) // down_y + 1
    out_w = (in_w * up_x + pad_x0 + pad_x1 - kernel_w) // down_x + 1

    return out.view(-1, channel, out_h, out_w)


================================================
FILE: op/upfirdn2d_kernel.cu
================================================
// 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

#include <torch/types.h>

#include <ATen/ATen.h>
#include <ATen/AccumulateType.h>
#include <ATen/cuda/CUDAApplyUtils.cuh>
#include <ATen/cuda/CUDAContext.h>

#include <cuda.h>
#include <cuda_runtime.h>

static __host__ __device__ __forceinline__ int floor_div(int a, int b) {
  int c = a / b;

  if (c * b > a) {
    c--;
  }

  return c;
}

struct UpFirDn2DKernelParams {
  int up_x;
  int up_y;
  int down_x;
  int down_y;
  int pad_x0;
  int pad_x1;
  int pad_y0;
  int pad_y1;

  int major_dim;
  int in_h;
  int in_w;
  int minor_dim;
  int kernel_h;
  int kernel_w;
  int out_h;
  int out_w;
  int loop_major;
  int loop_x;
};

template <typename scalar_t>
__global__ void upfirdn2d_kernel_large(scalar_t *out, const scalar_t *input,
                                       const scalar_t *kernel,
                                       const UpFirDn2DKernelParams p) {
  int minor_idx = blockIdx.x * blockDim.x + threadIdx.x;
  int out_y = minor_idx / p.minor_dim;
  minor_idx -= out_y * p.minor_dim;
  int out_x_base = blockIdx.y * p.loop_x * blockDim.y + threadIdx.y;
  int major_idx_base = blockIdx.z * p.loop_major;

  if (out_x_base >= p.out_w || out_y >= p.out_h ||
      major_idx_base >= p.major_dim) {
    return;
  }

  int mid_y = out_y * p.down_y + p.up_y - 1 - p.pad_y0;
  int in_y = min(max(floor_div(mid_y, p.up_y), 0), p.in_h);
  int h = min(max(floor_div(mid_y + p.kernel_h, p.up_y), 0), p.in_h) - in_y;
  int kernel_y = mid_y + p.kernel_h - (in_y + 1) * p.up_y;

  for (int loop_major = 0, major_idx = major_idx_base;
       loop_major < p.loop_major && major_idx < p.major_dim;
       loop_major++, major_idx++) {
    for (int loop_x = 0, out_x = out_x_base;
         loop_x < p.loop_x && out_x < p.out_w; loop_x++, out_x += blockDim.y) {
      int mid_x = out_x * p.down_x + p.up_x - 1 - p.pad_x0;
      int in_x = min(max(floor_div(mid_x, p.up_x), 0), p.in_w);
      int w = min(max(floor_div(mid_x + p.kernel_w, p.up_x), 0), p.in_w) - in_x;
      int kernel_x = mid_x + p.kernel_w - (in_x + 1) * p.up_x;

      const scalar_t *x_p =
          &input[((major_idx * p.in_h + in_y) * p.in_w + in_x) * p.minor_dim +
                 minor_idx];
      const scalar_t *k_p = &kernel[kernel_y * p.kernel_w + kernel_x];
      int x_px = p.minor_dim;
      int k_px = -p.up_x;
      int x_py = p.in_w * p.minor_dim;
      int k_py = -p.up_y * p.kernel_w;

      scalar_t v = 0.0f;

      for (int y = 0; y < h; y++) {
        for (int x = 0; x < w; x++) {
          v += static_cast<scalar_t>(*x_p) * static_cast<scalar_t>(*k_p);
          x_p += x_px;
          k_p += k_px;
        }

        x_p += x_py - w * x_px;
        k_p += k_py - w * k_px;
      }

      out[((major_idx * p.out_h + out_y) * p.out_w + out_x) * p.minor_dim +
          minor_idx] = v;
    }
  }
}

template <typename scalar_t, int up_x, int up_y, int down_x, int down_y,
          int kernel_h, int kernel_w, int tile_out_h, int tile_out_w>
__global__ void upfirdn2d_kernel(scalar_t *out, const scalar_t *input,
                                 const scalar_t *kernel,
                                 const UpFirDn2DKernelParams p) {
  const int tile_in_h = ((tile_out_h - 1) * down_y + kernel_h - 1) / up_y + 1;
  const int tile_in_w = ((tile_out_w - 1) * down_x + kernel_w - 1) / up_x + 1;

  __shared__ volatile float sk[kernel_h][kernel_w];
  __shared__ volatile float sx[tile_in_h][tile_in_w];

  int minor_idx = blockIdx.x;
  int tile_out_y = minor_idx / p.minor_dim;
  minor_idx -= tile_out_y * p.minor_dim;
  tile_out_y *= tile_out_h;
  int tile_out_x_base = blockIdx.y * p.loop_x * tile_out_w;
  int major_idx_base = blockIdx.z * p.loop_major;

  if (tile_out_x_base >= p.out_w | tile_out_y >= p.out_h |
      major_idx_base >= p.major_dim) {
    return;
  }

  for (int tap_idx = threadIdx.x; tap_idx < kernel_h * kernel_w;
       tap_idx += blockDim.x) {
    int ky = tap_idx / kernel_w;
    int kx = tap_idx - ky * kernel_w;
    scalar_t v = 0.0;

    if (kx < p.kernel_w & ky < p.kernel_h) {
      v = kernel[(p.kernel_h - 1 - ky) * p.kernel_w + (p.kernel_w - 1 - kx)];
    }

    sk[ky][kx] = v;
  }

  for (int loop_major = 0, major_idx = major_idx_base;
       loop_major < p.loop_major & major_idx < p.major_dim;
       loop_major++, major_idx++) {
    for (int loop_x = 0, tile_out_x = tile_out_x_base;
         loop_x < p.loop_x & tile_out_x < p.out_w;
         loop_x++, tile_out_x += tile_out_w) {
      int tile_mid_x = tile_out_x * down_x + up_x - 1 - p.pad_x0;
      int tile_mid_y = tile_out_y * down_y + up_y - 1 - p.pad_y0;
      int tile_in_x = floor_div(tile_mid_x, up_x);
      int tile_in_y = floor_div(tile_mid_y, up_y);

      __syncthreads();

      for (int in_idx = threadIdx.x; in_idx < tile_in_h * tile_in_w;
           in_idx += blockDim.x) {
        int rel_in_y = in_idx / tile_in_w;
        int rel_in_x = in_idx - rel_in_y * tile_in_w;
        int in_x = rel_in_x + tile_in_x;
        int in_y = rel_in_y + tile_in_y;

        scalar_t v = 0.0;

        if (in_x >= 0 & in_y >= 0 & in_x < p.in_w & in_y < p.in_h) {
          v = input[((major_idx * p.in_h + in_y) * p.in_w + in_x) *
                        p.minor_dim +
                    minor_idx];
        }

        sx[rel_in_y][rel_in_x] = v;
      }

      __syncthreads();
      for (int out_idx = threadIdx.x; out_idx < tile_out_h * tile_out_w;
           out_idx += blockDim.x) {
        int rel_out_y = out_idx / tile_out_w;
        int rel_out_x = out_idx - rel_out_y * tile_out_w;
        int out_x = rel_out_x + tile_out_x;
        int out_y = rel_out_y + tile_out_y;

        int mid_x = tile_mid_x + rel_out_x * down_x;
        int mid_y = tile_mid_y + rel_out_y * down_y;
        int in_x = floor_div(mid_x, up_x);
        int in_y = floor_div(mid_y, up_y);
        int rel_in_x = in_x - tile_in_x;
        int rel_in_y = in_y - tile_in_y;
        int kernel_x = (in_x + 1) * up_x - mid_x - 1;
        int kernel_y = (in_y + 1) * up_y - mid_y - 1;

        scalar_t v = 0.0;

#pragma unroll
        for (int y = 0; y < kernel_h / up_y; y++)
#pragma unroll
          for (int x = 0; x < kernel_w / up_x; x++)
            v += sx[rel_in_y + y][rel_in_x + x] *
                 sk[kernel_y + y * up_y][kernel_x + x * up_x];

        if (out_x < p.out_w & out_y < p.out_h) {
          out[((major_idx * p.out_h + out_y) * p.out_w + out_x) * p.minor_dim +
              minor_idx] = v;
        }
      }
    }
  }
}

torch::Tensor upfirdn2d_op(const torch::Tensor &input,
                           const torch::Tensor &kernel, int up_x, int up_y,
                           int down_x, int down_y, int pad_x0, int pad_x1,
                           int pad_y0, int pad_y1) {
  int curDevice = -1;
  cudaGetDevice(&curDevice);
  cudaStream_t stream = at::cuda::getCurrentCUDAStream(curDevice);

  UpFirDn2DKernelParams p;

  auto x = input.contiguous();
  auto k = kernel.contiguous();

  p.major_dim = x.size(0);
  p.in_h = x.size(1);
  p.in_w = x.size(2);
  p.minor_dim = x.size(3);
  p.kernel_h = k.size(0);
  p.kernel_w = k.size(1);
  p.up_x = up_x;
  p.up_y = up_y;
  p.down_x = down_x;
  p.down_y = down_y;
  p.pad_x0 = pad_x0;
  p.pad_x1 = pad_x1;
  p.pad_y0 = pad_y0;
  p.pad_y1 = pad_y1;

  p.out_h = (p.in_h * p.up_y + p.pad_y0 + p.pad_y1 - p.kernel_h + p.down_y) /
            p.down_y;
  p.out_w = (p.in_w * p.up_x + p.pad_x0 + p.pad_x1 - p.kernel_w + p.down_x) /
            p.down_x;

  auto out =
      at::empty({p.major_dim, p.out_h, p.out_w, p.minor_dim}, x.options());

  int mode = -1;

  int tile_out_h = -1;
  int tile_out_w = -1;

  if (p.up_x == 1 && p.up_y == 1 && p.down_x == 1 && p.down_y == 1 &&
      p.kernel_h <= 4 && p.kernel_w <= 4) {
    mode = 1;
    tile_out_h = 16;
    tile_out_w = 64;
  }

  if (p.up_x == 1 && p.up_y == 1 && p.down_x == 1 && p.down_y == 1 &&
      p.kernel_h <= 3 && p.kernel_w <= 3) {
    mode = 2;
    tile_out_h = 16;
    tile_out_w = 64;
  }

  if (p.up_x == 2 && p.up_y == 2 && p.down_x == 1 && p.down_y == 1 &&
      p.kernel_h <= 4 && p.kernel_w <= 4) {
    mode = 3;
    tile_out_h = 16;
    tile_out_w = 64;
  }

  if (p.up_x == 2 && p.up_y == 2 && p.down_x == 1 && p.down_y == 1 &&
      p.kernel_h <= 2 && p.kernel_w <= 2) {
    mode = 4;
    tile_out_h = 16;
    tile_out_w = 64;
  }

  if (p.up_x == 1 && p.up_y == 1 && p.down_x == 2 && p.down_y == 2 &&
      p.kernel_h <= 4 && p.kernel_w <= 4) {
    mode = 5;
    tile_out_h = 8;
    tile_out_w = 32;
  }

  if (p.up_x == 1 && p.up_y == 1 && p.down_x == 2 && p.down_y == 2 &&
      p.kernel_h <= 2 && p.kernel_w <= 2) {
    mode = 6;
    tile_out_h = 8;
    tile_out_w = 32;
  }

  dim3 block_size;
  dim3 grid_size;

  if (tile_out_h > 0 && tile_out_w > 0) {
    p.loop_major = (p.major_dim - 1) / 16384 + 1;
    p.loop_x = 1;
    block_size = dim3(32 * 8, 1, 1);
    grid_size = dim3(((p.out_h - 1) / tile_out_h + 1) * p.minor_dim,
                     (p.out_w - 1) / (p.loop_x * tile_out_w) + 1,
                     (p.major_dim - 1) / p.loop_major + 1);
  } else {
    p.loop_major = (p.major_dim - 1) / 16384 + 1;
    p.loop_x = 4;
    block_size = dim3(4, 32, 1);
    grid_size = dim3((p.out_h * p.minor_dim - 1) / block_size.x + 1,
                     (p.out_w - 1) / (p.loop_x * block_size.y) + 1,
                     (p.major_dim - 1) / p.loop_major + 1);
  }

  AT_DISPATCH_FLOATING_TYPES_AND_HALF(x.scalar_type(), "upfirdn2d_cuda", [&] {
    switch (mode) {
    case 1:
      upfirdn2d_kernel<scalar_t, 1, 1, 1, 1, 4, 4, 16, 64>
          <<<grid_size, block_size, 0, stream>>>(out.data_ptr<scalar_t>(),
                                                 x.data_ptr<scalar_t>(),
                                                 k.data_ptr<scalar_t>(), p);

      break;

    case 2:
      upfirdn2d_kernel<scalar_t, 1, 1, 1, 1, 3, 3, 16, 64>
          <<<grid_size, block_size, 0, stream>>>(out.data_ptr<scalar_t>(),
                                                 x.data_ptr<scalar_t>(),
                                                 k.data_ptr<scalar_t>(), p);

      break;

    case 3:
      upfirdn2d_kernel<scalar_t, 2, 2, 1, 1, 4, 4, 16, 64>
          <<<grid_size, block_size, 0, stream>>>(out.data_ptr<scalar_t>(),
                                                 x.data_ptr<scalar_t>(),
                                                 k.data_ptr<scalar_t>(), p);

      break;

    case 4:
      upfirdn2d_kernel<scalar_t, 2, 2, 1, 1, 2, 2, 16, 64>
          <<<grid_size, block_size, 0, stream>>>(out.data_ptr<scalar_t>(),
                                                 x.data_ptr<scalar_t>(),
                                                 k.data_ptr<scalar_t>(), p);

      break;

    case 5:
      upfirdn2d_kernel<scalar_t, 1, 1, 2, 2, 4, 4, 8, 32>
          <<<grid_size, block_size, 0, stream>>>(out.data_ptr<scalar_t>(),
                                                 x.data_ptr<scalar_t>(),
                                                 k.data_ptr<scalar_t>(), p);

      break;

    case 6:
      upfirdn2d_kernel<scalar_t, 1, 1, 2, 2, 4, 4, 8, 32>
          <<<grid_size, block_size, 0, stream>>>(out.data_ptr<scalar_t>(),
                                                 x.data_ptr<scalar_t>(),
                                                 k.data_ptr<scalar_t>(), p);

      break;

    default:
      upfirdn2d_kernel_large<scalar_t><<<grid_size, block_size, 0, stream>>>(
          out.data_ptr<scalar_t>(), x.data_ptr<scalar_t>(),
          k.data_ptr<scalar_t>(), p);
    }
  });

  return out;
}

================================================
FILE: prepare_data.py
================================================
import argparse
from io import BytesIO
import multiprocessing
from functools import partial

from PIL import Image

import lmdb
from tqdm import tqdm
from torchvision import datasets
from torchvision.transforms import functional as trans_fn

# ImageFile.LOAD_TRUNCATED_IMAGES = True


def resize_and_convert(img, size, resample, quality=100):
    img = trans_fn.resize(img, size, resample)
    img = trans_fn.center_crop(img, size)
    buffer = BytesIO()
    img.save(buffer, format="jpeg", quality=quality)
    val = buffer.getvalue()

    return val


def resize_multiple(img, sizes=(128, 256, 512, 1024), resample=Image.LANCZOS, quality=100):
    imgs = []

    for size in sizes:
        imgs.append(resize_and_convert(img, size, resample, quality))

    return imgs


def resize_worker(img_file, sizes, resample):
    i, file = img_file
    try:
        img = Image.open(file)
        img = img.convert("RGB")
    except:
        print(file, "truncated")
    out = resize_multiple(img, sizes=sizes, resample=resample)

    return i, out


def prepare(env, dataset, n_worker, sizes=(128, 256, 512, 1024), resample=Image.LANCZOS):
    resize_fn = partial(resize_worker, sizes=sizes, resample=resample)

    files = sorted(dataset.imgs, key=lambda x: x[0])
    files = [(i, file) for i, (file, label) in enumerate(files)]
    total = 0

    with multiprocessing.Pool(n_worker) as pool:
        for i, imgs in tqdm(pool.imap_unordered(resize_fn, files)):
            for size, img in zip(sizes, imgs):
                key = f"{size}-{str(i).zfill(5)}".encode("utf-8")

                with env.begin(write=True) as txn:
                    txn.put(key, img)

            total += 1

        with env.begin(write=True) as txn:
            txn.put("length".encode("utf-8"), str(total).encode("utf-8"))


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--out", type=str)
    parser.add_argument("--size", type=str, default="128,256,512,1024")
    parser.add_argument("--n_worker", type=int, default=8)
    parser.add_argument("--resample", type=str, default="bilinear")
    parser.add_argument("path", type=str)

    args = parser.parse_args()

    resample_map = {"lanczos": Image.LANCZOS, "bilinear": Image.BILINEAR}
    resample = resample_map[args.resample]

    sizes = [int(s.strip()) for s in args.size.split(",")]

    print(f"Make dataset of image sizes:", ", ".join(str(s) for s in sizes))

    imgset = datasets.ImageFolder(args.path)

    with lmdb.open(args.out, map_size=1024 ** 4, readahead=False) as env:
        prepare(env, imgset, args.n_worker, sizes=sizes, resample=resample)


================================================
FILE: prepare_vae_codes.py
================================================
import argparse
import numpy as np
import multiprocessing
from functools import partial

import lmdb
from tqdm import tqdm

import torch as th
from autoencoder import ConvSegNet
from torchvision import datasets
import torchvision.transforms as transforms


def lmdmb_write_worker(i_code, env, size):
    i, code = i_code.cpu().numpy()
    key = f"{size}-{str(i).zfill(5)}".encode("utf-8")
    with env.begin(write=True) as txn:
        txn.put(key, code)


def prepare(env, vae, loader, total, batch_size, n_worker, size=1024):
    write_fn = partial(lmdmb_write_worker, env=env, size=size)

    b = 0
    with multiprocessing.Pool(n_worker) as pool:
        for batch in tqdm(loader):
            code_nums = np.arange(b * batch_size, (b + 1) * batch_size)

            with th.no_grad():
                codes = vae.module.encode(batch[0].cuda())

            pool.imap_unordered(write_fn, zip(code_nums, codes))

            b += 1

    with env.begin(write=True) as txn:
        txn.put("length".encode("utf-8"), str(total).encode("utf-8"))


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--out", type=str)
    parser.add_argument("--size", type=int, default=1024)
    parser.add_argument("--n_worker", type=int, default=24)
    parser.add_argument("--batch_size", type=int, default=4)
    parser.add_argument("--resample", type=str, default="bilinear")
    parser.add_argument("data_path", type=str)
    parser.add_argument("vae_checkpoint", type=str)

    args = parser.parse_args()

    print(f"Make dataset of image size:", args.size)

    transform = transforms.Compose(
        [
            transforms.Resize(args.size),
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True),
        ]
    )
    imgset = datasets.ImageFolder(args.data_path, transform=transform)
    loader = th.utils.data.DataLoader(imgset, batch_size=args.batch_size, num_workers=int(args.n_worker / 2))
    print(args.batch_size)
    print(loader)

    vae = ConvSegNet()
    vae.load_state_dict(th.load(args.vae_checkpoint)["vae"])
    vae = th.nn.DataParallel(vae).eval().cuda()

    with lmdb.open(args.out, map_size=1024 ** 4, readahead=False) as env:
        prepare(
            env,
            vae,
            loader,
            total=len(imgset),
            batch_size=args.batch_size,
            n_worker=int(args.n_worker / 2),
            size=args.size,
        )


================================================
FILE: projector.py
================================================
import argparse
import math
import os

import torch
from torch import optim
from torch.nn import functional as F
from torchvision import transforms
from PIL import Image
from tqdm import tqdm

import lpips
from model import Generator


def noise_regularize(noises):
    loss = 0

    for noise in noises:
        size = noise.shape[2]

        while True:
            loss = (
                loss
                + (noise * torch.roll(noise, shifts=1, dims=3)).mean().pow(2)
                + (noise * torch.roll(noise, shifts=1, dims=2)).mean().pow(2)
            )

            if size <= 8:
                break

            noise = noise.reshape([1, 1, size // 2, 2, size // 2, 2])
            noise = noise.mean([3, 5])
            size //= 2

    return loss


def noise_normalize_(noises):
    for noise in noises:
        mean = noise.mean()
        std = noise.std()

        noise.data.add_(-mean).div_(std)


def get_lr(t, initial_lr, rampdown=0.25, rampup=0.05):
    lr_ramp = min(1, (1 - t) / rampdown)
    lr_ramp = 0.5 - 0.5 * math.cos(lr_ramp * math.pi)
    lr_ramp = lr_ramp * min(1, t / rampup)

    return initial_lr * lr_ramp


def latent_noise(latent, strength):
    noise = torch.randn_like(latent) * strength

    return latent + noise


def make_image(tensor):
    return (
        tensor.detach()
        .clamp_(min=-1, max=1)
        .add(1)
        .div_(2)
        .mul(255)
        .type(torch.uint8)
        .permute(0, 2, 3, 1)
        .to("cpu")
        .numpy()
    )


if __name__ == "__main__":
    device = "cuda"

    parser = argparse.ArgumentParser()
    parser.add_argument("--ckpt", type=str, required=True)
    parser.add_argument("--size", type=int, default=256)
    parser.add_argument("--lr_rampup", type=float, default=0.05)
    parser.add_argument("--lr_rampdown", type=float, default=0.25)
    parser.add_argument("--lr", type=float, default=0.1)
    parser.add_argument("--noise", type=float, default=0.05)
    parser.add_argument("--noise_ramp", type=float, default=0.75)
    parser.add_argument("--step", type=int, default=1000)
    parser.add_argument("--noise_regularize", type=float, default=1e5)
    parser.add_argument("--mse", type=float, default=0)
    parser.add_argument("--w_plus", action="store_true")
    parser.add_argument("files", metavar="FILES", nargs="+")

    args = parser.parse_args()

    n_mean_latent = 10000

    resize = min(args.size, 256)

    transform = transforms.Compose(
        [
            transforms.Resize(resize),
            transforms.CenterCrop(resize),
            transforms.ToTensor(),
            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
        ]
    )

    imgs = []

    for imgfile in args.files:
        img = transform(Image.open(imgfile).convert("RGB"))
        imgs.append(img)

    imgs = torch.stack(imgs, 0).to(device)

    g_ema = Generator(args.size, 512, 8)
    g_ema.load_state_dict(torch.load(args.ckpt)["g_ema"], strict=False)
    g_ema.eval()
    g_ema = g_ema.to(device)

    with torch.no_grad():
        noise_sample = torch.randn(n_mean_latent, 512, device=device)
        latent_out = g_ema.style(noise_sample)

        latent_mean = latent_out.mean(0)
        latent_std = ((latent_out - latent_mean).pow(2).sum() / n_mean_latent) ** 0.5

    percept = lpips.PerceptualLoss(model="net-lin", net="vgg", use_gpu=device.startswith("cuda"))

    noises = g_ema.make_noise()

    latent_in = latent_mean.detach().clone().unsqueeze(0).repeat(2, 1)

    if args.w_plus:
        latent_in = latent_in.unsqueeze(1).repeat(1, g_ema.n_latent, 1)

    latent_in.requires_grad = True

    for noise in noises:
        noise.requires_grad = True

    optimizer = optim.Adam([latent_in] + noises, lr=args.lr)

    pbar = tqdm(range(args.step))
    latent_path = []

    for i in pbar:
        t = i / args.step
        lr = get_lr(t, args.lr)
        optimizer.param_groups[0]["lr"] = lr
        noise_strength = latent_std * args.noise * max(0, 1 - t / args.noise_ramp) ** 2
        latent_n = latent_noise(latent_in, noise_strength.item())

        img_gen, _ = g_ema([latent_n], input_is_latent=True, noise=noises)

        batch, channel, height, width = img_gen.shape

        if height > 256:
            factor = height // 256

            img_gen = img_gen.reshape(batch, channel, height // factor, factor, width // factor, factor)
            img_gen = img_gen.mean([3, 5])

        p_loss = percept(img_gen, imgs).sum()
        n_loss = noise_regularize(noises)
        mse_loss = F.mse_loss(img_gen, imgs)

        loss = p_loss + args.noise_regularize * n_loss + args.mse * mse_loss

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        noise_normalize_(noises)

        if (i + 1) % 100 == 0:
            latent_path.append(latent_in.detach().clone())

        pbar.set_description(
            (
                f"perceptual: {p_loss.item():.4f}; noise regularize: {n_loss.item():.4f};"
                f" mse: {mse_loss.item():.4f}; lr: {lr:.4f}"
            )
        )

    result_file = {"noises": noises}

    img_gen, _ = g_ema([latent_path[-1]], input_is_latent=True, noise=noises)

    filename = os.path.splitext(os.path.basename(args.files[0]))[0] + ".pt"

    img_ar = make_image(img_gen)

    for i, input_name in enumerate(args.files):
        result_file[input_name] = {"img": img_gen[i], "latent": latent_in[i]}
        img_name = os.path.splitext(os.path.basename(input_name))[0] + "-project.png"
        pil_img = Image.fromarray(img_ar[i])
        pil_img.save(img_name)

    torch.save(result_file, filename)


================================================
FILE: render.py
================================================
import queue
from threading import Thread

import ffmpeg
import numpy as np
import PIL.Image
import torch as th
from tqdm import tqdm

th.set_grad_enabled(False)
th.backends.cudnn.benchmark = True


def render(
    generator,
    latents,
    noise,
    offset,
    duration,
    batch_size,
    out_size,
    output_file,
    audio_file=None,
    truncation=1.0,
    bends=[],
    rewrites={},
    randomize_noise=False,
    ffmpeg_preset="slow",
):
    split_queue = queue.Queue()
    render_queue = queue.Queue()

    # postprocesses batched torch tensors to individual RGB numpy arrays
    def split_batches(jobs_in, jobs_out):
        while True:
            try:
                imgs = jobs_in.get(timeout=5)
            except queue.Empty:
                return
            imgs = (imgs.clamp_(-1, 1) + 1) * 127.5
            imgs = imgs.permute(0, 2, 3, 1)
            for img in imgs:
                jobs_out.put(img.cpu().numpy().astype(np.uint8))
            jobs_in.task_done()

    # start background ffmpeg process that listens on stdin for frame data
    if out_size == 512:
        output_size = "512x512"
    elif out_size == 1024:
        output_size = "1024x1024"
    elif out_size == 1920:
        output_size = "1920x1080"
    elif out_size == 1080:
        output_size = "1080x1920"
    else:
        raise Exception("The only output sizes currently supported are: 512, 1024, 1080, or 1920")

    if audio_file is not None:
        audio = ffmpeg.input(audio_file, ss=offset, t=duration, guess_layout_max=0)
        video = (
            ffmpeg.input("pipe:", format="rawvideo", pix_fmt="rgb24", framerate=len(latents) / duration, s=output_size)
            .output(
                audio,
                output_file,
                framerate=len(latents) / duration,
                vcodec="libx264",
                pix_fmt="yuv420p",
                preset=ffmpeg_preset,
                audio_bitrate="320K",
                ac=2,
                v="warning",
            )
            .global_args("-hide_banner")
            .overwrite_output()
            .run_async(pipe_stdin=True)
        )
    else:
        video = (
            ffmpeg.input("pipe:", format="rawvideo", pix_fmt="rgb24", framerate=len(latents) / duration, s=output_size)
            .output(
                output_file,
                framerate=len(latents) / duration,
                vcodec="libx264",
                pix_fmt="yuv420p",
                preset=ffmpeg_preset,
                v="warning",
            )
            .global_args("-hide_banner")
            .overwrite_output()
            .run_async(pipe_stdin=True)
        )

    # writes numpy frames to ffmpeg stdin as raw rgb24 bytes
    def make_video(jobs_in):
        w, h = [int(dim) for dim in output_size.split("x")]
        for _ in tqdm(range(len(latents)), position=0, leave=True, ncols=80):
            img = jobs_in.get(timeout=5)
            if img.shape[1] == 2048:
                img = img[:, 112:-112, :]
                im = PIL.Image.fromarray(img)
                img = np.array(im.resize((1920, 1080), PIL.Image.BILINEAR))
            elif img.shape[0] == 2048:
                img = img[112:-112, :, :]
                im = PIL.Image.fromarray(img)
                img = np.array(im.resize((1080, 1920), PIL.Image.BILINEAR))
            assert (
                img.shape[1] == w and img.shape[0] == h
            ), f"""generator's output image size does not match specified output size: \n
                got: {img.shape[1]}x{img.shape[0]}\t\tshould be {output_size}"""
            video.stdin.write(img.tobytes())
            jobs_in.task_done()
        video.stdin.close()
        video.wait()

    splitter = Thread(target=split_batches, args=(split_queue, render_queue))
    splitter.daemon = True
    renderer = Thread(target=make_video, args=(render_queue,))
    renderer.daemon = True

    # make all data that needs to be loaded to the GPU float, contiguous, and pinned
    # the entire process is severly memory-transfer bound, but at least this might help a little
    latents = latents.float().contiguous().pin_memory()

    for ni, noise_scale in enumerate(noise):
        noise[ni] = noise_scale.float().contiguous().pin_memory() if noise_scale is not None else None

    param_dict = dict(generator.named_parameters())
    original_weights = {}
    for param, (rewrite, modulation) in rewrites.items():
        rewrites[param] = [rewrite, modulation.float().contiguous().pin_memory()]
        original_weights[param] = param_dict[param].copy().cpu().float().contiguous().pin_memory()

    for bend in bends:
        if "modulation" in bend:
            bend["modulation"] = bend["modulation"].float().contiguous().pin_memory()

    if not isinstance(truncation, float):
        truncation = truncation.float().contiguous().pin_memory()

    for n in range(0, len(latents), batch_size):
        # load batches of data onto the GPU
        latent_batch = latents[n : n + batch_size].cuda(non_blocking=True)

        noise_batch = []
        for noise_scale in noise:
            if noise_scale is not None:
                noise_batch.append(noise_scale[n : n + batch_size].cuda(non_blocking=True))
            else:
                noise_batch.append(None)

        bend_batch = []
        if bends is not None:
            for bend in bends:
                if "modulation" in bend:
                    transform = bend["transform"](bend["modulation"][n : n + batch_size].cuda(non_blocking=True))
                    bend_batch.append({"layer": bend["layer"], "transform": transform})
                else:
                    bend_batch.append({"layer": bend["layer"], "transform": bend["transform"]})

        for param, (rewrite, modulation) in rewrites.items():
            transform = rewrite(modulation[n : n + batch_size])
            rewritten_weight = transform(original_weights[param]).cuda(non_blocking=True)
            param_attrs = param.split(".")
            mod = generator
            for attr in param_attrs[:-1]:
                mod = getattr(mod, attr)
            setattr(mod, param_attrs[-1], th.nn.Parameter(rewritten_weight))

        if not isinstance(truncation, float):
            truncation_batch = truncation[n : n + batch_size].cuda(non_blocking=True)
        else:
            truncation_batch = truncation

        # forward through the generator
        outputs, _ = generator(
            styles=latent_batch,
            noise=noise_batch,
            truncation=truncation_batch,
            transform_dict_list=bend_batch,
            randomize_noise=randomize_noise,
            input_is_latent=True,
        )

        # send output to be split into frames and rendered one by one
        split_queue.put(outputs)

        if n == 0:
            splitter.start()
            renderer.start()

    splitter.join()
    renderer.join()


def write_video(arr, output_file, fps):
    print(f"writing {arr.shape[0]} frames...")

    output_size = "x".join(reversed([str(s) for s in arr.shape[1:-1]]))

    ffmpeg_proc = (
        ffmpeg.input("pipe:", format="rawvideo", pix_fmt="rgb24", framerate=fps, s=output_size)
        .output(output_file, framerate=fps, vcodec="libx264", preset="slow", v="warning")
        .global_args("-benchmark", "-stats", "-hide_banner")
        .overwrite_output()
        .run_async(pipe_stdin=True)
    )

    for frame in arr:
        ffmpeg_proc.stdin.write(frame.astype(np.uint8).tobytes())

    ffmpeg_proc.stdin.close()
    ffmpeg_proc.wait()


================================================
FILE: requirements.txt
================================================
torch
torchvision
numpy
librosa
cython
madmom
tqdm
kornia
matplotlib
ffmpeg-python
joblib

================================================
FILE: select_latents.py
================================================
import gc, math
import argparse
import tkinter as tk
import numpy as np
from PIL import Image, ImageTk
import torch as th
import torch.nn.functional as F
import torchvision
from models.stylegan2 import Generator as G_style2
import tkinter as tk

# --- classes ---
try:
    from Tkinter import Canvas, Frame
    from ttk import Scrollbar

    from Tkconstants import *
except ImportError:
    from tkinter import Canvas, Frame
    from tkinter.ttk import Scrollbar

    from tkinter.constants import *

import platform

OS = platform.system()


class HoverButton(tk.Button):
    def __init__(self, master, **kw):
        tk.Button.__init__(self, master=master, **kw)
        self.defaultBackground = self["background"]
        self.bind("<Enter>", self.on_enter)
        self.bind("<Leave>", self.on_leave)

    def on_enter(self, e):
        self["background"] = self["activebackground"]

    def on_leave(self, e):
        self["background"] = self.defaultBackground


class InvisibleScrollbar(Scrollbar):
    def set(self, lo, hi):
        self.tk.call("grid", "remove", self)
        Scrollbar.set(self, lo, hi)


class Mousewheel_Support(object):
    # implemetation of singleton pattern
    _instance = None

    def __new__(cls, *args, **kwargs):
        if not cls._instance:
            cls._instance = object.__new__(cls)
        return cls._instance

    def __init__(self, root, horizontal_factor=1, vertical_factor=1):
        self._active_area = None
        if isinstance(horizontal_factor, int):
            self.horizontal_factor = horizontal_factor
        else:
            raise Exception("Vertical factor must be an integer.")
        if isinstance(vertical_factor, int):
            self.vertical_factor = vertical_factor
        else:
            raise Exception("Horizontal factor must be an integer.")
        if OS == "Linux":
            root.bind_all("<4>", self._on_mousewheel, add="+")
            root.bind_all("<5>", self._on_mousewheel, add="+")
        else:
            # Windows and MacOS
            root.bind_all("<MouseWheel>", self._on_mousewheel, add="+")

    def _on_mousewheel(self, event):
        if self._active_area:
            self._active_area.onMouseWheel(event)

    def _mousewheel_bind(self, widget):
        self._active_area = widget

    def _mousewheel_unbind(self):
        self._active_area = None

    def add_support_to(
        self, widget=None, xscrollbar=None, yscrollbar=None, what="units", horizontal_factor=None, vertical_factor=None
    ):
        if xscrollbar is None and yscrollbar is None:
            return
        if xscrollbar is not None:
            horizontal_factor = horizontal_factor or self.horizontal_factor
            xscrollbar.onMouseWheel = self._make_mouse_wheel_handler(widget, "x", self.horizontal_factor, what)
            xscrollbar.bind("<Enter>", lambda event, scrollbar=xscrollbar: self._mousewheel_bind(scrollbar))
            xscrollbar.bind("<Leave>", lambda event: self._mousewheel_unbind())
        if yscrollbar is not None:
            vertical_factor = vertical_factor or self.vertical_factor
            yscrollbar.onMouseWheel = self._make_mouse_wheel_handler(widget, "y", self.vertical_factor, what)
            yscrollbar.bind("<Enter>", lambda event, scrollbar=yscrollbar: self._mousewheel_bind(scrollbar))
            yscrollbar.bind("<Leave>", lambda event: self._mousewheel_unbind())
        main_scrollbar = yscrollbar if yscrollbar is not None else xscrollbar
        if widget is not None:
            if isinstance(widget, list) or isinstance(widget, tuple):
                list_of_widgets = widget
                for widget in list_of_widgets:
                    widget.bind("<Enter>", lambda event: self._mousewheel_bind(widget))
                    widget.bind("<Leave>", lambda event: self._mousewheel_unbind())
                    widget.onMouseWheel = main_scrollbar.onMouseWheel
            else:
                widget.bind("<Enter>", lambda event: self._mousewheel_bind(widget))
                widget.bind("<Leave>", lambda event: self._mousewheel_unbind())
                widget.onMouseWheel = main_scrollbar.onMouseWheel

    @staticmethod
    def _make_mouse_wheel_handler(widget, orient, factor=1 / 120, what="units"):
        view_command = getattr(widget, orient + "view")
        if OS == "Linux":

            def onMouseWheel(event):
                if event.num == 4:
                    view_command("scroll", (-1) * factor, what)
                elif event.num == 5:
                    view_command("scroll", factor, what)

        elif OS == "Windows":

            def onMouseWheel(event):
                view_command("scroll", (-1) * int((event.delta / 120) * factor), what)

        elif OS == "Darwin":

            def onMouseWheel(event):
                view_command("scroll", event.delta, what)

        return onMouseWheel


class Scrolling_Area(Frame, object):
    def __init__(
        self,
        master,
        width=None,
        anchor=N,
        height=None,
        mousewheel_speed=2,
        scroll_horizontally=True,
        xscrollbar=None,
        scroll_vertically=True,
        yscrollbar=None,
        background="black",
        inner_frame=Frame,
        **kw,
    ):
        Frame.__init__(self, master, class_="Scrolling_Area", background=background)
        self.grid_columnconfigure(0, weight=1)
        self.grid_rowconfigure(0, weight=1)
        self._width = width
        self._height = height
        self.canvas = Canvas(self, background=background, highlightthickness=0, width=width, height=height)
        self.canvas.grid(row=0, column=0, sticky=N + E + W + S)
        if scroll_vertically:
            if yscrollbar is not None:
                self.yscrollbar = yscrollbar
            else:
                self.yscrollbar = InvisibleScrollbar(self, orient=VERTICAL)
                self.yscrollbar.grid(row=0, column=1, sticky=N + S)
            self.canvas.configure(yscrollcommand=self.yscrollbar.set)
            self.yscrollbar["command"] = self.canvas.yview
        else:
            self.yscrollbar = None
        if scroll_horizontally:
            if xscrollbar is not None:
                self.xscrollbar = xscrollbar
            else:
                self.xscrollbar = InvisibleScrollbar(self, orient=HORIZONTAL)
                self.xscrollbar.grid(row=1, column=0, sticky=E + W)
            self.canvas.configure(xscrollcommand=self.xscrollbar.set)
            self.xscrollbar["command"] = self.canvas.xview
        else:
            self.xscrollbar = None
        self.rowconfigure(0, weight=1)
        self.columnconfigure(0, weight=1)
        self.innerframe = inner_frame(self.canvas, **kw)
        self.innerframe.pack(anchor=anchor)
        self.canvas.create_window(0, 0, window=self.innerframe, anchor="nw", tags="inner_frame")
        self.canvas.bind("<Configure>", self._on_canvas_configure)
        Mousewheel_Support(self).add_support_to(self.canvas, xscrollbar=self.xscrollbar, yscrollbar=self.yscrollbar)

    @property
    def width(self):
        return self.canvas.winfo_width()

    @width.setter
    def width(self, width):
        self.canvas.configure(width=width)

    @property
    def height(self):
        return self.canvas.winfo_height()

    @height.setter
    def height(self, height):
        self.canvas.configure(height=height)

    def set_size(self, width, height):
        self.canvas.configure(width=width, height=height)

    def _on_canvas_configure(self, event):
        width = max(self.innerframe.winfo_reqwidth(), event.width)
        height = max(self.innerframe.winfo_reqheight(), event.height)
        self.canvas.configure(scrollregion="0 0 %s %s" % (width, height))
        self.canvas.itemconfigure("inner_frame", width=width, height=height)

    def update_viewport(self):
        self.update()
        window_width = self.innerframe.winfo_reqwidth()
        window_height = self.innerframe.winfo_reqheight()
        if self._width is None:
            canvas_width = window_width
        else:
            canvas_width = min(self._width, window_width)
        if self._height is None:
            canvas_height = window_height
        else:
            canvas_height = min(self._height, window_height)
        self.canvas.configure(
            scrollregion="0 0 %s %s" % (window_width, window_height), width=self._width, height=self._height
        )
        self.canvas.itemconfigure("inner_frame", width=window_width, height=window_height)


th.set_grad_enabled(False)
th.backends.cudnn.benchmark = True

parser = argparse.ArgumentParser()

parser.add_argument("--ckpt", type=str)
parser.add_argument("--res", type=int, default=1024)
parser.add_argument("--output_dir", type=str, default="./workspace/")
parser.add_argument("--truncation", type=float, default=1.5)
parser.add_argument("--noconst", action="store_false")

args = parser.parse_args()

name = args.ckpt.split("/")[-1].split(".")[0]
GENERATOR = (
    G_style2(size=args.res, style_dim=512, n_mlp=8, checkpoint=args.ckpt, output_size=1024, constant_input=args.noconst)
    .eval()
    .cuda()
)
# GENERATOR = G_style(checkpoint=args.ckpt, output_size=1024).eval().cuda()
# GENERATOR = th.nn.DataParallel(GENERATOR)

IMAGES_PER_ROW = 4
IMSIZE = (1920 - 240) // IMAGES_PER_ROW

ALL_LATENTS = []
DROP_IDXS = []
INTRO_IDXS = []
IMAGES = []


def generate_images(n):
    imgs = []
    for _ in range(n // 8):
        random_latents = th.randn(8, 512).cuda()
        mapped_latents = GENERATOR(random_latents, noise=None, truncation=args.truncation, map_latents=True)
        for latent in mapped_latents:
            ALL_LATENTS.append(latent[None, ...].cpu().numpy())
        batch, _ = GENERATOR(
            styles=mapped_latents,
            noise=None,
            truncation=args.truncation,
            transform_dict_list=[],
            randomize_noise=True,
            input_is_latent=True,
        )
        imgs.append(batch)
    imgs = th.cat(imgs)[:n]
    imgs = F.interpolate(imgs, IMSIZE, mode="bilinear", align_corners=False)
    imgs = (imgs.clamp_(-1, 1) + 1) * 127.5
    imgs = imgs.permute(0, 2, 3, 1)
    imgs_np = imgs.cpu().numpy().astype(np.uint8)
    del imgs
    gc.collect()
    th.cuda.empty_cache()
    return imgs_np


root = tk.Tk()
root.title(name)

imgrid = Scrolling_Area(root, bg="black", width=1680, height=1080)
imgrid.pack(side="left", expand=True, fill="both")

panel = tk.Frame(root, relief="flat", bg="black")
panel.pack(side="right", expand=True, fill="both")


def render_latents(latents):
    imgs = []
    for i in range(latents.shape[0] // 8 + 1):
        if len(latents[8 * i : 8 * (i + 1)]) < 1:
            continue
        batch, _ = GENERATOR(
            styles=latents[8 * i : 8 * (i + 1)].cuda(),
            noise=None,
            truncation=args.truncation,
            transform_dict_list=[],
            randomize_noise=True,
            input_is_latent=True,
        )
        imgs.append(batch)
    imgs = th.cat(imgs)
    imgs = (imgs.clamp_(-1, 1) + 1) / 2
    return imgs


def save():
    intro_latents = np.concatenate(ALL_LATENTS)[INTRO_IDXS]
    torchvision.utils.save_image(
        render_latents(th.from_numpy(intro_latents)),
        f"{args.output_dir}/{name}_intro_latents.jpg",
        nrow=int(round(math.sqrt(intro_latents.shape[0]) * 4 / 3)),
        padding=0,
        normalize=False,
    )
    np.save(f"{args.output_dir}/{name}_intro_latents.npy", intro_latents)

    drop_latents = np.concatenate(ALL_LATENTS)[DROP_IDXS]
    torchvision.utils.save_image(
        render_latents(th.from_numpy(drop_latents)),
        f"{args.output_dir}/{name}_drop_latents.jpg",
        nrow=int(round(math.sqrt(drop_latents.shape[0]) * 4 / 3)),
        padding=0,
        normalize=False,
    )
    np.save(f"{args.output_dir}/{name}_drop_latents.npy", drop_latents)


tk.Label(panel, text="latents", height=3, bg="black", fg="white").pack(side="top")

but = HoverButton(
    panel,
    text="Save",
    command=save,
    height=3,
    width=8,
    bg="black",
    fg="white",
    activebackground="#333333",
    activeforeground="white",
    relief="flat",
    highlightbackground="#333333",
)
but.pack(side="bottom")

intro = tk.LabelFrame(panel, text="intro", width=240, height=490, bg="black", fg="white", relief="flat")
intro.pack(side="top", fill="both")

drop = tk.LabelFrame(panel, text="drop", width=240, height=490, bg="black", fg="white", relief="flat")
drop.pack(side="bottom", fill="both")

introgrid = Scrolling_Area(intro, width=240, height=490, bg="black")
introgrid.pack(side="top", fill="both")

dropgrid = Scrolling_Area(drop, width=240, height=490, bg="black")
dropgrid.pack(side="bottom", fill="both")

im_num = 0
intro_im_num = 0
drop_im_num = 0


def add_intro(label):
    global intro_im_num
    img_id = int(label.__str__().split(".")[-1])

    INTRO_IDXS.append(img_id)

    img = ImageTk.PhotoImage(image=IMAGES[img_id].resize((46, 46), Image.ANTIALIAS))
    lbl = tk.Label(introgrid.innerframe, image=img, borderwidth=0, highlightthickness=0)
    lbl.image = img  # this line need to prevent gc
    lbl.grid(row=math.floor(intro_im_num / 5), column=intro_im_num % 5)
    lbl.bind("<Button-1>", lambda event, l=lbl: remove_intro(l))

    intro_im_num += 1
    introgrid.update_viewport()


def remove_intro(label):
    remove_idx = list(reversed(introgrid.innerframe.grid_slaves())).index(label)
    label.grid_remove()
    del INTRO_IDXS[remove_idx]
    global intro_im_num
    intro_im_num = 0
    for im in reversed(introgrid.innerframe.grid_slaves()):
        im.grid_remove()
        im.grid(row=math.floor(intro_im_num / 5), column=intro_im_num % 5)
        intro_im_num += 1
    introgrid.update_viewport()


def add_drop(label):
    global drop_im_num
    img_id = int(label.__str__().split(".")[-1])

    DROP_IDXS.append(img_id)

    img = ImageTk.PhotoImage(image=IMAGES[img_id].resize((46, 46), Image.ANTIALIAS))
    lbl = tk.Label(dropgrid.innerframe, image=img, borderwidth=0, highlightthickness=0)
    lbl.image = img  # this line need to prevent gc
    lbl.grid(row=math.floor(drop_im_num / 5), column=drop_im_num % 5)
    lbl.bind("<Button-1>", lambda event, l=lbl: remove_drop(l))

    drop_im_num += 1
    dropgrid.update_viewport()


def remove_drop(label):
    remove_idx = list(reversed(dropgrid.innerframe.grid_slaves())).index(label)
    label.grid_remove()
    del DROP_IDXS[remove_idx]
    global drop_im_num
    drop_im_num = 0
    for im in reversed(dropgrid.innerframe.grid_slaves()):
        im.grid_remove()
        im.grid(row=math.floor(drop_im_num / 5), column=drop_im_num % 5)
        drop_im_num += 1
    dropgrid.update_viewport()


def add_images(n):
    global im_num, IMAGES

    for im_arr in generate_images(n):
        im = Image.fromarray(im_arr)
        IMAGES.append(im)
        img = ImageTk.PhotoImage(image=im)

        label = tk.Label(imgrid.innerframe, image=img, name=str(im_num), borderwidth=0, highlightthickness=0)
        label.image = img  # this line need to prevent gc
        label.grid(row=math.floor(im_num / IMAGES_PER_ROW), column=im_num % IMAGES_PER_ROW)
        label.bind("<Button-1>", lambda event, l=label: add_intro(l))
        label.bind("<Button-3>", lambda event, l=label: add_drop(l))

        im_num += 1

    HoverButton(
        imgrid.innerframe,
        text="More",
        command=lambda n=35: add_images(n),
        height=3,
        width=8,
        bg="black",
        fg="white",
        activebackground="#333333",
        activeforeground="white",
        relief="flat",
        highlightbackground="#333333",
    ).grid(
        row=math.floor(im_num / IMAGES_PER_ROW) + 1,
        column=math.floor(IMAGES_PER_ROW / 2 - 1),
        columnspan=1 if math.floor(IMAGES_PER_ROW / 2 - 1) % 2 == 0 else 2,
    )
    imgrid.update_viewport()


add_images(24)


root.mainloop()


================================================
FILE: train.py
================================================
import argparse
import gc
import math
import os
import random
import sys
import time

import numpy as np
import torch as th
import wandb
from contrastive_learner import ContrastiveLearner, RandomApply
from kornia import augmentation as augs
from scipy.ndimage import gaussian_filter
from torch import nn
from torch.nn import functional as F
from torch.utils import data
from torchvision import transforms, utils
from tqdm import tqdm

import validation
from augment import augment
from dataset import MultiResolutionDataset
from distributed import get_rank, reduce_loss_dict, reduce_sum, synchronize
from lookahead_minimax import LookaheadMinimax
from models.stylegan2 import Discriminator, Generator

sys.path.insert(0, "../lookahead_minimax")


def data_sampler(dataset, shuffle, distributed):
    if distributed:
        return data.distributed.DistributedSampler(dataset, shuffle=shuffle)
    if shuffle:
        return data.RandomSampler(dataset)
    else:
        return data.SequentialSampler(dataset)


def requires_grad(model, flag=True):
    for p in model.parameters():
        p.requires_grad = flag


def accumulate(model1, model2, decay=0.5 ** (32.0 / 10_000)):
    par1 = dict(model1.named_parameters())
    par2 = dict(model2.named_parameters())
    for name, param in model1.named_parameters():
        param.data = decay * par1[name].data + (1 - decay) * par2[name].data


def sample_data(loader):
    while True:
        for batch in loader:
            yield batch


def make_noise(batch_size, latent_dim, prob):
    if prob > 0 and random.random() < prob:
        return th.randn(2, batch_size, latent_dim, device=device).unbind(0)
    else:
        return [th.randn(batch_size, latent_dim, device=device)]


def d_logistic_loss(real_pred, fake_pred):
    real_loss = F.softplus(-real_pred)
    fake_loss = F.softplus(fake_pred)
    return real_loss.mean() + fake_loss.mean()


def d_r1_penalty(real_img, real_pred, args):
    (grad_real,) = th.autograd.grad(real_pred.sum(), real_img, create_graph=True)
    r1_loss = grad_real.pow(2).view(grad_real.shape[0], -1).sum(1).mean()
    r1_loss = r1_loss / 2.0 + 0 * real_pred[0]
    return r1_loss


def g_non_saturating_loss(fake_pred):
    return F.softplus(-fake_pred).mean()


def g_path_length_regularization(generator, mean_path_length, args):
    path_batch_size = max(1, args.batch_size // args.path_batch_shrink)

    noise = make_noise(path_batch_size, args.latent_size, args.mixing_prob)
    fake_img, latents = generator(noise, return_latents=True)

    img_noise = th.randn_like(fake_img) / math.sqrt(fake_img.shape[2] * fake_img.shape[3])
    noisy_img_sum = (fake_img * img_noise).sum()

    (grad,) = th.autograd.grad(noisy_img_sum, latents, create_graph=True)

    path_lengths = th.sqrt(grad.pow(2).sum(2).mean(1))
    path_mean = mean_path_length + 0.01 * (path_lengths.mean() - mean_path_length)
    path_loss = (path_lengths - path_mean).pow(2).mean()
    if not th.isnan(path_mean):
        mean_path_length = path_mean.detach()

    if args.path_batch_shrink:
        path_loss += 0 * fake_img[0, 0, 0, 0]

    return path_loss, mean_path_length


def train(args, loader, generator, discriminator, contrast_learner, g_optim, d_optim, g_ema):
    if args.distributed:
        g_module = generator.module
        d_module = discriminator.module
        if contrast_learner is not None:
            cl_module = contrast_learner.module
    else:
        g_module = generator
        d_module = discriminator
        cl_module = contrast_learner

    loader = sample_data(loader)
    sample_z = th.randn(args.n_sample, args.latent_size, device=device)
    mse = th.nn.MSELoss()
    mean_path_length = th.cuda.FloatTensor([0.0])
    ada_aug_signs = th.cuda.FloatTensor([0.0])
    ada_aug_n = th.cuda.FloatTensor([0.0])
    ada_aug_p = th.cuda.FloatTensor([args.augment_p if args.augment_p > 0 else 0.0])
    ada_aug_step = th.cuda.FloatTensor([args.ada_target / args.ada_length])
    r_t_stat = th.cuda.FloatTensor([0.0])
    fids = []

    loss_dict = {
        "Generator": th.cuda.FloatTensor([0.0]),
        "Discriminator": th.cuda.FloatTensor([0.0]),
        "Real Score": th.cuda.FloatTensor([0.0]),
        "Fake Score": th.cuda.FloatTensor([0.0]),
        "Contrastive": th.cuda.FloatTensor([0.0]),
        "Consistency": th.cuda.FloatTensor([0.0]),
        "R1 Penalty": th.cuda.FloatTensor([0.0]),
        "Path Length Regularization": th.cuda.FloatTensor([0.0]),
        "Augment": th.cuda.FloatTensor([0.0]),
        "Rt": th.cuda.FloatTensor([0.0]),
    }

    pbar = range(args.iter)
    if get_rank() == 0:
        pbar = tqdm(pbar, initial=args.start_iter, dynamic_ncols=True, smoothing=0)
    for idx in pbar:
        i = idx + args.start_iter
        if i > args.iter:
            print("Done!")
            break
        tick_start = time.time()

        for k, v in loss_dict.items():
            loss_dict[k].mul_(0)

        requires_grad(generator, False)
        requires_grad(discriminator, True)

        discriminator.zero_grad()
        for _ in range(args.num_accumulate):
            real_img_og = next(loader).to(device, non_blocking=True)

            noise = make_noise(args.batch_size, args.latent_size, args.mixing_prob)
            fake_img_og, _ = generator(noise)
            if args.augment:
                fake_img, _ = augment(fake_img_og, ada_aug_p)
                real_img, _ = augment(real_img_og, ada_aug_p)
            else:
                fake_img = fake_img_og
                real_img = real_img_og

            fake_pred = discriminator(fake_img)
            real_pred = discriminator(real_img)
            logistic_loss = d_logistic_loss(real_pred, fake_pred)
            loss_dict["Discriminator"] += logistic_loss.detach()
            loss_dict["Real Score"] += real_pred.mean().detach()
            loss_dict["Fake Score"] += fake_pred.mean().detach()
            d_loss = logistic_loss

            if args.contrastive > 0:
                contrast_learner(fake_img_og, fake_img, accumulate=True)
                contrast_learner(real_img_og, real_img, accumulate=True)
                contrast_loss = cl_module.calculate_loss()
                loss_dict["Contrastive"] += contrast_loss.detach()
                d_loss += args.contrastive * contrast_loss

            if args.balanced_consistency > 0:
                consistency_loss = mse(real_pred, discriminator(real_img_og)) + mse(
                    fake_pred, discriminator(fake_img_og)
                )
                loss_dict["Consistency"] += consistency_loss.detach()
                d_loss += args.balanced_consistency * consistency_loss

            d_loss /= args.num_accumulate
            d_loss.backward()
        d_optim.step()

        if args.r1 > 0 and i % args.d_reg_every == 0:
            discriminator.zero_grad()
            for _ in range(args.num_accumulate):
                real_img = next(loader).to(device, non_blocking=True)
                real_img.requires_grad = True
                real_pred = discriminator(real_img)
                r1_loss = d_r1_penalty(real_img, real_pred, args)
                loss_dict["R1 Penalty"] += r1_loss.detach().squeeze()
                r1_loss = args.r1 * args.d_reg_every * r1_loss / args.num_accumulate
                r1_loss.backward()
            d_optim.step()

        if args.augment and args.augment_p == 0:
            ada_aug_signs += th.sign(real_pred).sum().item()
            ada_aug_n += real_pred.shape[0]
            ada_aug_signs, ada_aug_n = reduce_sum(ada_aug_signs), reduce_sum(ada_aug_n)

            if ada_aug_n > 255:
                r_t_stat = ada_aug_signs / ada_aug_n
                loss_dict["Rt"] += r_t_stat
                if r_t_stat > args.ada_target:
                    sign = 1
                else:
                    sign = -1

                ada_aug_p += sign * ada_aug_step * ada_aug_n
                ada_aug_p = th.clamp(ada_aug_p, 0, 1)
                ada_aug_signs.mul_(0)
                ada_aug_n.mul_(0)
                loss_dict["Augment"] += ada_aug_p

        requires_grad(generator, True)
        requires_grad(discriminator, False)

        generator.zero_grad()
        for _ in range(args.num_accumulate):
            noise = make_noise(args.batch_size, args.latent_size, args.mixing_prob)
            fake_img, _ = generator(noise)
            if args.augment:
                fake_img, _ = augment(fake_img, ada_aug_p)
            fake_pred = discriminator(fake_img)
            g_loss = g_non_saturating_loss(fake_pred)
            loss_dict["Generator"] += g_loss.detach()
            g_loss /= args.num_accumulate
            g_loss.backward()
        g_optim.step()

        if args.path_regularize > 0 and i % args.g_reg_every == 0:
            generator.zero_grad()
            for _ in range(args.num_accumulate):
                path_loss, mean_path_length = g_path_length_regularization(generator, mean_path_length, args)
                loss_dict["Path Length Regularization"] += path_loss.detach()
                path_loss = args.path_regularize * args.g_reg_every * path_loss / args.num_accumulate
                path_loss.backward()
            g_optim.step()

        accumulate(g_ema, g_module)

        loss_reduced = reduce_loss_dict(loss_dict)
        log_dict = {k: v.mean().item() / args.num_accumulate for k, v in loss_reduced.items() if v != 0}
        log_dict["Tick Length"] = time.time() - tick_start

        if get_rank() == 0:
            with th.no_grad():
                if args.log_spec_norm:
                    G_norms = []
                    for name, spec_norm in g_module.named_buffers():
                        if "spectral_norm" in name:
                            G_norms.append(spec_norm.cpu().numpy())
                    G_norms = np.array(G_norms)
                    D_norms = []
                    for name, spec_norm in d_module.named_buffers():
                        if "spectral_norm" in name:
                            D_norms.append(spec_norm.cpu().numpy())
                    D_norms = np.array(D_norms)
                    log_dict[f"Spectral Norms/G min spectral norm"] = np.log(G_norms).min()
                    log_dict[f"Spectral Norms/G mean spectral norm"] = np.log(G_norms).mean()
                    log_dict[f"Spectral Norms/G max spectral norm"] = np.log(G_norms).max()
                    log_dict[f"Spectral Norms/D min spectral norm"] = np.log(D_norms).min()
                    log_dict[f"Spectral Norms/D mean spectral norm"] = np.log(D_norms).mean()
                    log_dict[f"Spectral Norms/D max spectral norm"] = np.log(D_norms).max()

                if args.img_every != -1 and i % args.img_every == 0:
                    g_ema.eval()
                    sample = []
                    for sub in range(0, len(sample_z), args.batch_size):
                        subsample, _ = g_ema([sample_z[sub : sub + args.batch_size]])
                        sample.append(subsample.detach().cpu())
                    sample = th.cat(sample).detach()
                    grid = utils.make_grid(sample, nrow=10, normalize=True, range=(-1, 1))
                    log_dict["Generated Images EMA"] = [wandb.Image(grid, caption=f"Step {i}")]

                if args.eval_every != -1 and i % args.eval_every == 0:
                    fid_dict = validation.fid(
                        g_ema, args.val_batch_size, args.fid_n_sample, args.fid_truncation, args.name
                    )

                    fid = fid_dict["FID"]
                    fids.append(fid)
                    density = fid_dict["Density"]
                    coverage = fid_dict["Coverage"]

                    ppl = validation.ppl(
                        g_ema, args.val_batch_size, args.ppl_n_sample, args.ppl_space, args.ppl_crop, args.latent_size,
                    )

                    log_dict["Evaluation/FID"] = fid
                    log_dict["Sweep/FID_smooth"] = gaussian_filter(np.array(fids), [5])[-1]
                    log_dict["Evaluation/Density"] = density
                    log_dict["Evaluation/Coverage"] = coverage
                    log_dict["Evaluation/PPL"] = ppl

                wandb.log(log_dict)

                if args.eval_every != -1:
                    description = (
                        f"FID: {fid:.4f}   PPL: {ppl:.4f}   Dens: {density:.4f}   Cov: {coverage:.4f}   "
                        + f"G: {log_dict['Generator']:.4f}   D: {log_dict['Discriminator']:.4f}"
                    )
                else:
                    description = f"G: {log_dict['Generator']:.4f}   D: {log_dict['Discriminator']:.4f}"
                if "Augment" in log_dict:
                    description += f"   Aug: {log_dict['Augment']:.4f}"  #   Rt: {log_dict['Rt']:.4f}"
                if "R1 Penalty" in log_dict:
                    description += f"   R1: {log_dict['R1 Penalty']:.4f}"
                if "Path Length Regularization" in log_dict:
                    description += f"   Path: {log_dict['Path Length Regularization']:.4f}"
                pbar.set_description(description)

                if i % args.checkpoint_every == 0:
                    check_name = "-".join(
                        [
                            args.name,
                            args.wbname,
                            wandb.run.dir.split("/")[-1].split("-")[-1],
                            # str(int(fid)),
                            str(args.size),
                            str(i).zfill(6),
                        ]
                    )
                    th.save(
                        {
                            "g": g_module.state_dict(),
                            "d": d_module.state_dict(),
                            # "cl": cl_module.state_dict(),
                            "g_ema": g_ema.state_dict(),
                            "g_optim": g_optim.state_dict(),
                            "d_optim": d_optim.state_dict(),
                        },
                        f"/home/hans/modelzoo/maua-sg2/{check_name}.pt",
                    )

        if args.profile_mem:
            gpu_profile(frame=sys._getframe(), event="line", arg=None)


if __name__ == "__main__":
    device = "cuda"

    parser = argparse.ArgumentParser()

    parser.add_argument("--wbname", type=str, required=True)
    parser.add_argument("--wbproj", type=str, required=True)
    parser.add_argument("--wbgroup", type=str, default=None)

    # data options
    parser.add_argument("--path", type=str, required=True)
    parser.add_argument("--vflip", type=bool, default=False)
    parser.add_argument("--hflip", type=bool, default=True)

    # training options
    parser.add_argument("--batch_size", type=int, default=12)
    parser.add_argument("--num_accumulate", type=int, default=1)

    parser.add_argument("--checkpoint", type=str, default=None)
    parser.add_argument("--transfer_mapping_only", type=bool, default=False)
    parser.add_argument("--start_iter", type=int, default=0)
    parser.add_argument("--iter", type=int, default=20_000)

    # model options
    parser.add_argument("--size", type=int, default=1024)
    parser.add_argument("--min_rgb_size", type=int, default=4)
    parser.add_argument("--latent_size", type=int, default=512)
    parser.add_argument("--n_mlp", type=int, default=8)
    parser.add_argument("--n_sample", type=int, default=60)
    parser.add_argument("--constant_input", type=bool, default=False)
    parser.add_argument("--channel_multiplier", type=int, default=2)
    parser.add_argument("--d_skip", type=bool, default=True)

    # optimizer options
    parser.add_argument("--lr", type=float, default=0.002)
    parser.add_argument("--d_lr_ratio", type=float, default=1.0)
    parser.add_argument("--lookahead", type=bool, default=True)
    parser.add_argument("--la_steps", type=float, default=500)
    parser.add_argument("--la_alpha", type=float, default=0.5)

    # loss options
    parser.add_argument("--r1", type=float, default=1e-5)
    parser.add_argument("--path_regularize", type=float, default=1)
    parser.add_argument("--path_batch_shrink", type=int, default=2)
    parser.add_argument("--d_reg_every", type=int, default=16)
    parser.add_argument("--g_reg_every", type=int, default=4)
    parser.add_argument("--mixing_prob", type=float, default=0.666)

    # augmentation options
    parser.add_argument("--augment", type=bool, default=True)
    parser.add_argument("--contrastive", type=float, default=0)
    parser.add_argument("--balanced_consistency", type=float, default=0)
    parser.add_argument("--augment_p", type=float, default=0)
    parser.add_argument("--ada_target", type=float, default=0.6)
    parser.add_argument("--ada_length", type=int, default=15_000)

    # validation options
    parser.add_argument("--val_batch_size", type=int, default=6)
    parser.add_argument("--fid_n_sample", type=int, default=2500)
    parser.add_argument("--fid_truncation", type=float, default=None)
    parser.add_argument("--ppl_space", choices=["z", "w"], default="w")
    parser.add_argument("--ppl_n_sample", type=int, default=1250)
    parser.add_argument("--ppl_crop", type=bool, default=False)

    # logging options
    parser.add_argument("--log_spec_norm", type=bool, default=False)
    parser.add_argument("--img_every", type=int, default=1000)
    parser.add_argument("--eval_every", type=int, default=-1)
    parser.add_argument("--checkpoint_every", type=int, default=1000)
    parser.add_argument("--profile_mem", action="store_true")

    # (multi-)GPU options
    parser.add_argument("--local_rank", type=int, default=0)
    parser.add_argument("--cudnn_benchmark", type=bool, default=True)

    args = parser.parse_args()
    if args.balanced_consistency > 0 or args.contrastive > 0:
        args.augment = True
    args.name = os.path.splitext(os.path.basename(args.path))[0]
    args.r1 = args.r1 * args.size ** 2

    args.num_gpus = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1
    th.backends.cudnn.benchmark = args.cudnn_benchmark
    args.distributed = args.num_gpus > 1

    # code for updating wandb configs that were incorrect
    # if args.local_rank == 0:
    #     api = wandb.Api()
    #     run = api.run("wav/temperatuur/7kp6g0zt")
    #     run.config = vars(args)
    #     run.update()
    # exit()

    if args.distributed:
        th.cuda.set_device(args.local_rank)
        th.distributed.init_process_group(backend="nccl", init_method="env://")
        synchronize()

    generator = Generator(
        args.size,
        args.latent_size,
        args.n_mlp,
        channel_multiplier=args.channel_multiplier,
        constant_input=args.constant_input,
        min_rgb_size=args.min_rgb_size,
    ).to(device, non_blocking=True)
    discriminator = Discriminator(args.size, channel_multiplier=args.channel_multiplier, use_skip=args.d_skip).to(
        device
    )

    if args.log_spec_norm:
        for name, parameter in generator.named_parameters():
            if "weight" in name and parameter.squeeze().dim() > 1:
                mod = generator
                for attr in name.replace(".weight", "").split("."):
                    mod = getattr(mod, attr)
                validation.track_spectral_norm(mod)
        for name, parameter in discriminator.named_parameters():
            if "weight" in name and parameter.squeeze().dim() > 1:
                mod = discriminator
                for attr in name.replace(".weight", "").split("."):
                    mod = getattr(mod, attr)
                validation.track_spectral_norm(mod)

    g_ema = Generator(
        args.size,
        args.latent_size,
        args.n_mlp,
        channel_multiplier=args.channel_multiplier,
        constant_input=args.constant_input,
        min_rgb_size=args.min_rgb_size,
    ).to(device, non_blocking=True)
    g_ema.requires_grad_(False)
    g_ema.eval()
    accumulate(g_ema, generator, 0)

    if args.contrastive > 0:
        contrast_learner = ContrastiveLearner(
            discriminator,
            args.size,
            augment_fn=nn.Sequential(
                nn.ReflectionPad2d(int((math.sqrt(2) - 1) * args.size / 4)),  # zoom out
                augs.RandomHorizontalFlip(),
                RandomApply(augs.RandomAffine(degrees=0, translate=(0.25, 0.25), shear=(15, 15)), p=0.1),
                RandomApply(augs.RandomRotation(180), p=0.1),
                augs.RandomResizedCrop(size=(args.size, args.size), scale=(1, 1), ratio=(1, 1)),
                RandomApply(augs.RandomResizedCrop(size=(args.size, args.size), scale=(0.5, 0.9)), p=0.1),  # zoom in
                RandomApply(augs.RandomErasing(), p=0.1),
            ),
            hidden_layer=(-1, 0),
        )
    else:
        contrast_learner = None

    g_reg_ratio = args.g_reg_every / (args.g_reg_every + 1)
    d_reg_ratio = args.d_reg_every / (args.d_reg_every + 1)

    g_optim = th.optim.Adam(
        generator.parameters(), lr=args.lr * g_reg_ratio, betas=(0 ** g_reg_ratio, 0.99 ** g_reg_ratio),
    )
    d_optim = th.optim.Adam(
        discriminator.parameters(),
        lr=args.lr * d_reg_ratio * args.d_lr_ratio,
        betas=(0 ** d_reg_ratio, 0.99 ** d_reg_ratio),
    )

    if args.lookahead:
        g_optim = LookaheadMinimax(
            g_optim, d_optim, la_steps=args.la_steps, la_alpha=args.la_alpha, accumulate=args.num_accumulate
        )

    if args.checkpoint is not None:
        print("load model:", args.checkpoint)

        checkpoint = th.load(args.checkpoint, map_location=lambda storage, loc: storage)

        try:
            ckpt_name = os.path.basename(args.checkpoint)
            args.start_iter = int(os.path.splitext(ckpt_name)[-1].replace(args.name, ""))
        except ValueError:
            pass

        if args.transfer_mapping_only:
            print("Using generator latent mapping network from checkpoint")
            mapping_state_dict = {}
            for key, val in checkpoint["state_dict"].items():
                if "generator.style" in key:
                    mapping_state_dict[key.replace("generator.", "")] = val
            generator.load_state_dict(mapping_state_dict, strict=False)
        else:
            generator.load_state_dict(checkpoint["g"], strict=False)
            g_ema.load_state_dict(checkpoint["g_ema"], strict=False)

            discriminator.load_state_dict(checkpoint["d"], strict=False)

            if args.lookahead:
                g_optim.load_state_dict(checkpoint["g_optim"], checkpoint["d_optim"])
            else:
                g_optim.load_state_dict(checkpoint["g_optim"])
                d_optim.load_state_dict(checkpoint["d_optim"])

        del checkpoint
        th.cuda.empty_cache()

    if args.distributed:
        generator = nn.parallel.DistributedDataParallel(
            generator,
            device_ids=[args.local_rank],
            output_device=args.local_rank,
            broadcast_buffers=False,
            find_unused_parameters=True,
        )

        discriminator = nn.parallel.DistributedDataParallel(
            discriminator,
            device_ids=[args.local_rank],
            output_device=args.local_rank,
            broadcast_buffers=False,
            find_unused_parameters=True,
        )

        if contrast_learner is not None:
            contrast_learner = nn.parallel.DistributedDataParallel(
                contrast_learner,
                device_ids=[args.local_rank],
                output_device=args.local_rank,
                broadcast_buffers=False,
                find_unused_parameters=True,
            )

    transform = transforms.Compose(
        [
            transforms.RandomVerticalFlip(p=0.5 if args.vflip else 0),
            transforms.RandomHorizontalFlip(p=0.5 if args.hflip else 0),
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
        ]
    )

    dataset = MultiResolutionDataset(args.path, transform, args.size)
    loader = data.DataLoader(
        dataset,
        batch_size=args.batch_size,
        sampler=data_sampler(dataset, shuffle=True, distributed=args.distributed),
        num_workers=8,
        drop_last=True,
        pin_memory=True,
    )

    if get_rank() == 0:
        validation.get_dataset_inception_features(loader, args.name, args.size)
        if args.wbgroup is None:
            wandb.init(project=args.wbproj, name=args.wbname, config=vars(args))
        else:
            wandb.init(project=args.wbproj, group=args.wbgroup, name=args.wbname, config=vars(args))

    if args.profile_mem:
        os.environ["GPU_DEBUG"] = str(args.local_rank)
        from gpu_profile import gpu_profile

        sys.settrace(gpu_profile)

    train(args, loader, generator, discriminator, contrast_learner, g_optim, d_optim, g_ema)


================================================
FILE: train_profile.py
================================================
import argparse
import gc
import math
import os
import random
import sys
import time

import numpy as np
import torch as th
import torch.autograd.profiler as profiler
import wandb
from contrastive_learner import ContrastiveLearner, RandomApply
from kornia import augmentation as augs
from scipy.ndimage import gaussian_filter
from torch import nn
from torch.nn import functional as F
from torch.utils import data
from torchvision import transforms, utils
from tqdm import tqdm

import validation
from augment import augment
from dataset import MultiResolutionDataset
from distributed import get_rank, reduce_loss_dict, reduce_sum, synchronize
from lookahead_minimax import LookaheadMinimax
from models.stylegan2 import Discriminator, Generator

sys.path.insert(0, "../lookahead_minimax")


def data_sampler(dataset, shuffle, distributed):
    if distributed:
        return data.distributed.DistributedSampler(dataset, shuffle=shuffle)
    if shuffle:
        return data.RandomSampler(dataset)
    else:
        return data.SequentialSampler(dataset)


def requires_grad(model, flag=True):
    for p in model.parameters():
        p.requires_grad = flag


def accumulate(model1, model2, decay=0.5 ** (32.0 / 10_000)):
    par1 = dict(model1.named_parameters())
    par2 = dict(model2.named_parameters())
    for name, param in model1.named_parameters():
        param.data = decay * par1[name].data + (1 - decay) * par2[name].data


def sample_data(loader):
    while True:
        for batch in loader:
            yield batch


def make_noise(batch_size, latent_dim, prob):
    if prob > 0 and random.random() < prob:
        return th.randn(2, batch_size, latent_dim, device=device).unbind(0)
    else:
        return [th.randn(batch_size, latent_dim, device=device)]


def d_logistic_loss(real_pred, fake_pred):
    real_loss = F.softplus(-real_pred)
    fake_loss = F.softplus(fake_pred)
    return real_loss.mean() + fake_loss.mean()


def d_r1_penalty(real_img, real_pred, args):
    (grad_real,) = th.autograd.grad(real_pred.sum(), real_img, create_graph=True)
    r1_loss = grad_real.pow(2).view(grad_real.shape[0], -1).sum(1).mean()
    r1_loss = r1_loss / 2.0 + 0 * real_pred[0]
    return r1_loss


def g_non_saturating_loss(fake_pred):
    return F.softplus(-fake_pred).mean()


def g_path_length_regularization(generator, mean_path_length, args):
    path_batch_size = max(1, args.batch_size // args.path_batch_shrink)

    noise = make_noise(path_batch_size, args.latent_size, args.mixing_prob)
    fake_img, latents = generator(noise, return_latents=True)

    img_noise = th.randn_like(fake_img) / math.sqrt(fake_img.shape[2] * fake_img.shape[3])
    noisy_img_sum = (fake_img * img_noise).sum()

    (grad,) = th.autograd.grad(noisy_img_sum, latents, create_graph=True)

    path_lengths = th.sqrt(grad.pow(2).sum(2).mean(1))
    path_mean = mean_path_length + 0.01 * (path_lengths.mean() - mean_path_length)
    path_loss = (path_lengths - path_mean).pow(2).mean()
    if not th.isnan(path_mean):
        mean_path_length = path_mean.detach()

    if args.path_batch_shrink:
        path_loss += 0 * fake_img[0, 0, 0, 0]

    return path_loss, mean_path_length


# detach / item all the things
# separate out into smaller functions so those local scopes get cleaned better
# placeholder tensor pattern
# torch no grad where possible
def train(args, loader, generator, discriminator, contrast_learner, g_optim, d_optim, g_ema):
    if args.distributed:
        g_module = generator.module
        d_module = discriminator.module
        if contrast_learner is not None:
            cl_module = contrast_learner.module
    else:
        g_module = generator
        d_module = discriminator
        cl_module = contrast_learner

    loader = sample_data(loader)
    sample_z = th.randn(args.n_sample, args.latent_size, device=device)
    mse = th.nn.MSELoss()
    mean_path_length = 0
    ada_augment = th.tensor([0.0, 0.0], device=device)
    ada_aug_p = args.augment_p if args.augment_p > 0 else 0.0
    ada_aug_step = args.ada_target / args.ada_length
    r_t_stat = 0
    fids = []

    pbar = range(args.iter)
    if get_rank() == 0:
        pbar = tqdm(pbar, initial=args.start_iter, dynamic_ncols=True, smoothing=0)

    for idx in pbar:
        i = idx + args.start_iter
        if i > args.iter:
            print("Done!")
            break

        if idx == args.nsys_iter:
            print("Profiling begun at iteration {}".format(idx))
            th.cuda.cudart().cudaProfilerStart()
        if idx >= args.nsys_iter:
            th.cuda.nvtx.range_push("Iter {}".format(idx))

        tick_start = time.time()

        loss_dict = {
            "Generator": th.tensor(0, device=device).float(),
            "Discriminator": th.tensor(0, device=device).float(),
            "Real Score": th.tensor(0, device=device).float(),
            "Fake Score": th.tensor(0, device=device).float(),
            "Contrastive": th.tensor(0, device=device).float(),
            "Consistency": th.tensor(0, device=device).float(),
            "R1 Penalty": th.tensor(0, device=device).float(),
            "Path Length Regularization": th.tensor(0, device=device).float(),
            "Augment": th.tensor(0, device=device).float(),
            "Rt": th.tensor(0, device=device).float(),
        }

        with profiler.record_function("D train"):
            if idx >= args.nsys_iter:
                th.cuda.nvtx.range_push("D train")
            requires_grad(generator, False)
            requires_grad(discriminator, True)

            discriminator.zero_grad()
            for _ in range(args.num_accumulate):
                real_img_og = next(loader).to(device, non_blocking=True)
                noise = make_noise(args.batch_size, args.latent_size, args.mixing_prob)
                fake_img_og, _ = generator(noise)
                if args.augment:
                    fake_img, _ = augment(fake_img_og, ada_aug_p)
                    real_img, _ = augment(real_img_og, ada_aug_p)
                else:
                    fake_img = fake_img_og
                    real_img = real_img_og

                fake_pred = discriminator(fake_img)
                real_pred = discriminator(real_img)
                logistic_loss = d_logistic_loss(real_pred, fake_pred)
                loss_dict["Discriminator"] += logistic_loss.detach()
                loss_dict["Real Score"] += real_pred.mean().detach()
                loss_dict["Fake Score"] += fake_pred.mean().detach()
                d_loss = logistic_loss

                if args.contrastive > 0:
                    contrast_learner(fake_img_og, fake_img, accumulate=True)
                    contrast_learner(real_img_og, real_img, accumulate=True)
                    contrast_loss = cl_module.calculate_loss()
                    loss_dict["Contrastive"] += contrast_loss.detach()
                    d_loss += args.contrastive * contrast_loss

                if args.balanced_consistency > 0:
                    consistency_loss = mse(real_pred, discriminator(real_img_og)) + mse(
                        fake_pred, discriminator(fake_img_og)
                    )
                    loss_dict["Consistency"] += consistency_loss.detach()
                    d_loss += args.balanced_consistency * consistency_loss

                d_loss /= args.num_accumulate
                d_loss.backward()
            d_optim.step()
            if idx >= args.nsys_iter:
                th.cuda.nvtx.range_pop()

        if args.r1 > 0 and i % args.d_reg_every == 0:
            with profiler.record_function("D reg"):
                if idx >= args.nsys_iter:
                    th.cuda.nvtx.range_push("D reg")
                discriminator.zero_grad()
                for _ in range(args.num_accumulate):
                    real_img = next(loader).to(device, non_blocking=True)
                    real_img.requires_grad = True
                    real_pred = discriminator(real_img)
                    r1_loss = d_r1_penalty(real_img, real_pred, args)
                    loss_dict["R1 Penalty"] += r1_loss.detach().squeeze()
                    r1_loss = args.r1 * args.d_reg_every * r1_loss / args.num_accumulate
                    r1_loss.backward()
                d_optim.step()
                if idx >= args.nsys_iter:
                    th.cuda.nvtx.range_pop()

        if args.augment and args.augment_p == 0:
            with profiler.record_function("ADA"):
                if idx >= args.nsys_iter:
                    th.cuda.nvtx.range_push("ADA")
                ada_augment += th.tensor((th.sign(real_pred).sum().item(), real_pred.shape[0]), device=device)
                ada_augment = reduce_sum(ada_augment)

                if ada_augment[1] > 255:
                    pred_signs, n_pred = ada_augment.tolist()

                    r_t_stat = pred_signs / n_pred
                    loss_dict["Rt"] = th.tensor(r_t_stat, device=device).float()
                    if r_t_stat > args.ada_target:
                        sign = 1
                    else:
                        sign = -1

                    ada_aug_p += sign * ada_aug_step * n_pred
                    ada_aug_p = min(1, max(0, ada_aug_p))
                    ada_augment.mul_(0)
                    loss_dict["Augment"] = th.tensor(ada_aug_p, device=device).float()
                if idx >= args.nsys_iter:
                    th.cuda.nvtx.range_pop()

        with profiler.record_function("G train"):
            if idx >= args.nsys_iter:
                th.cuda.nvtx.range_push("G train")
            requires_grad(generator, True)
            requires_grad(discriminator, False)

            generator.zero_grad()
            for _ in range(args.num_accumulate):
                noise = make_noise(args.batch_size, args.latent_size, args.mixing_prob)
                fake_img, _ = generator(noise)
                if args.augment:
                    fake_img, _ = augment(fake_img, ada_aug_p)
                fake_pred = discriminator(fake_img)
                g_loss = g_non_saturating_loss(fake_pred)
                loss_dict["Generator"] += g_loss.detach()
                g_loss /= args.num_accumulate
                g_loss.backward()
            g_optim.step()
            if idx >= args.nsys_iter:
                th.cuda.nvtx.range_pop()

        if args.path_regularize > 0 and i % args.g_reg_every == 0:
            with profiler.record_function("G reg"):
                if idx >= args.nsys_iter:
                    th.cuda.nvtx.range_push("G reg")
                generator.zero_grad()
                for _ in range(args.num_accumulate):
                    path_loss, mean_path_length = g_path_length_regularization(generator, mean_path_length, args)
                    loss_dict["Path Length Regularization"] += path_loss.detach()
                    path_loss = args.path_regularize * args.g_reg_every * path_loss / args.num_accumulate
                    path_loss.backward()
                g_optim.step()
                if idx >= args.nsys_iter:
                    th.cuda.nvtx.range_pop()

        with profiler.record_function("Log / Eval"):
            if idx >= args.nsys_iter:
                th.cuda.nvtx.range_push("G accum")
            accumulate(g_ema, g_module)
            if idx >= args.nsys_iter:
                th.cuda.nvtx.range_pop()

            if idx >= args.nsys_iter:
                th.cuda.nvtx.range_push("Log / Eval")
            loss_reduced = reduce_loss_dict(loss_dict)
            log_dict = {k: v.mean().item() / args.num_accumulate for k, v in loss_reduced.items() if v != 0}
            log_dict["Tick Length"] = time.time() - tick_start

            if get_rank() == 0:
                with th.no_grad():
                    if args.log_spec_norm:
                        G_norms = []
                        for name, spec_norm in g_module.named_buffers():
                            if "spectral_norm" in name:
                                G_norms.append(spec_norm.cpu().numpy())
                        G_norms = np.array(G_norms)
                        D_norms = []
                        for name, spec_norm in d_module.named_buffers():
                            if "spectral_norm" in name:
                                D_norms.append(spec_norm.cpu().numpy())
                        D_norms = np.array(D_norms)
                        log_dict[f"Spectral Norms/G min spectral norm"] = np.log(G_norms).min()
                        log_dict[f"Spectral Norms/G mean spectral norm"] = np.log(G_norms).mean()
                        log_dict[f"Spectral Norms/G max spectral norm"] = np.log(G_norms).max()
                        log_dict[f"Spectral Norms/D min spectral norm"] = np.log(D_norms).min()
                        log_dict[f"Spectral Norms/D mean spectral norm"] = np.log(D_norms).mean()
                        log_dict[f"Spectral Norms/D max spectral norm"] = np.log(D_norms).max()

                    if args.img_every != -1 and i % args.img_every == 0:
                        g_ema.eval()
                        sample = []
                        for sub in range(0, len(sample_z), args.batch_size):
                            subsample, _ = g_ema([sample_z[sub : sub + args.batch_size]])
                            sample.append(subsample.detach().cpu())
                        sample = th.cat(sample)
                        grid = utils.make_grid(sample, nrow=10, normalize=True, range=(-1, 1))
                        log_dict["Generated Images EMA"] = [wandb.Image(grid, caption=f"Step {i}")]

                    if args.eval_every != -1 and i % args.eval_every == 0:
                        fid_dict = validation.fid(
                            g_ema, args.val_batch_size, args.fid_n_sample, args.fid_truncation, args.name
                        )

                        fid = fid_dict["FID"]
                        fids.append(fid)
                        density = fid_dict["Density"]
                        coverage = fid_dict["Coverage"]

                        ppl = validation.ppl(
                            g_ema,
                            args.val_batch_size,
                            args.ppl_n_sample,
                            args.ppl_space,
                            args.ppl_crop,
                            args.latent_size,
                        )

                        log_dict["Evaluation/FID"] = fid
                        log_dict["Sweep/FID_smooth"] = gaussian_filter(np.array(fids), [5])[-1]
                        log_dict["Evaluation/Density"] = density
                        log_dict["Evaluation/Coverage"] = coverage
                        log_dict["Evaluation/PPL"] = ppl

                        gc.collect()
                        th.cuda.empty_cache()

                    wandb.log(log_dict)

                    if args.eval_every != -1:
                        description = (
                            f"FID: {fid:.4f}   PPL: {ppl:.4f}   Dens: {density:.4f}   Cov: {coverage:.4f}   "
                            + f"G: {log_dict['Generator']:.4f}   D: {log_dict['Discriminator']:.4f}"
                        )
                    else:
                        description = f"G: {log_dict['Generator']:.4f}   D: {log_dict['Discriminator']:.4f}"
                    if "Augment" in log_dict:
                        description += f"   Aug: {log_dict['Augment']:.4f}"  #   Rt: {log_dict['Rt']:.4f}"
                    if "R1 Penalty" in log_dict:
                        description += f"   R1: {log_dict['R1 Penalty']:.4f}"
                    if "Path Length Regularization" in log_dict:
                        description += f"   Path: {log_dict['Path Length Regularization']:.4f}"
                    pbar.set_description(description)

                    if i % args.checkpoint_every == 0:
                        check_name = "-".join(
                            [
                                args.name,
                                args.wbname,
                                wandb.run.dir.split("/")[-1].split("-")[-1],
                                # str(int(fid)),
                                str(args.size),
                                str(i).zfill(6),
                            ]
                        )
                        th.save(
                            {
                                "g": g_module.state_dict(),
                                "d": d_module.state_dict(),
                                # "cl": cl_module.state_dict(),
                                "g_ema": g_ema.state_dict(),
                                "g_optim": g_optim.state_dict(),
                                "d_optim": d_optim.state_dict(),
                            },
                            f"/home/hans/modelzoo/maua-sg2/{check_name}.pt",
                        )

            if idx >= args.nsys_iter:
                th.cuda.nvtx.range_pop()
            if idx >= args.nsys_iter:
                th.cuda.nvtx.range_pop()  # iteration range

        gpu_profile(frame=sys._getframe(), event="line", arg=None)

    if args.nsys_iter != -1:
        th.cuda.cudart().cudaProfilerStop()


if __name__ == "__main__":
    device = "cuda"

    parser = argparse.ArgumentParser()

    parser.add_argument("--wbname", type=str, required=True)
    parser.add_argument("--wbproj", type=str, required=True)
    parser.add_argument("--wbgroup", type=str, default=None)

    # data options
    parser.add_argument("--path", type=str, required=True)
    parser.add_argument("--vflip", type=bool, default=False)
    parser.add_argument("--hflip", type=bool, default=True)

    # training options
    parser.add_argument("--batch_size", type=int, default=12)
    parser.add_argument("--num_accumulate", type=int, default=1)

    parser.add_argument("--checkpoint", type=str, default=None)
    parser.add_argument("--transfer_mapping_only", type=bool, default=False)
    parser.add_argument("--start_iter", type=int, default=0)
    parser.add_argument("--iter", type=int, default=60_000)

    # model options
    parser.add_argument("--size", type=int, default=256)
    parser.add_argument("--min_rgb_size", type=int, default=4)
    parser.add_argument("--latent_size", type=int, default=512)
    parser.add_argument("--n_mlp", type=int, default=8)
    parser.add_argument("--n_sample", type=int, default=60)
    parser.add_argument("--constant_input", type=bool, default=False)
    parser.add_argument("--channel_multiplier", type=int, default=2)
    parser.add_argument("--d_skip", type=bool, default=True)

    # optimizer options
    parser.add_argument("--lr", type=float, default=0.002)
    parser.add_argument("--d_lr_ratio", type=float, default=1.0)
    parser.add_argument("--lookahead", type=bool, default=True)
    parser.add_argument("--la_steps", type=float, default=500)
    parser.add_argument("--la_alpha", type=float, default=0.5)

    # loss options
    parser.add_argument("--r1", type=float, default=1e-5)
    parser.add_argument("--path_regularize", type=float, default=1)
    parser.add_argument("--path_batch_shrink", type=int, default=2)
    parser.add_argument("--d_reg_every", type=int, default=16)
    parser.add_argument("--g_reg_every", type=int, default=4)
    parser.add_argument("--mixing_prob", type=float, default=0.666)

    # augmentation options
    parser.add_argument("--augment", type=bool, default=False)
    parser.add_argument("--contrastive", type=float, default=0)
    parser.add_argument("--balanced_consistency", type=float, default=0)
    parser.add_argument("--augment_p", type=float, default=0)
    parser.add_argument("--ada_target", type=float, default=0.6)
    parser.add_argument("--ada_length", type=int, default=40_000)

    # validation options
    parser.add_argument("--val_batch_size", type=int, default=6)
    parser.add_argument("--fid_n_sample", type=int, default=2500)
    parser.add_argument("--fid_truncation", type=float, default=None)
    parser.add_argument("--ppl_space", choices=["z", "w"], default="w")
    parser.add_argument("--ppl_n_sample", type=int, default=1250)
    parser.add_argument("--ppl_crop", type=bool, default=False)

    # logging options
    parser.add_argument("--log_spec_norm", type=bool, default=False)
    parser.add_argument("--img_every", type=int, default=1000)
    parser.add_argument("--eval_every", type=int, default=1000)
    parser.add_argument("--checkpoint_every", type=int, default=1000)

    # (multi-)GPU options
    parser.add_argument("--local_rank", type=int, default=0)
    parser.add_argument("--cudnn_benchmark", type=bool, default=True)
    parser.add_argument("--nsys_iter", type=int, default=-1)
    parser.add_argument("--th_prof", action="store_true")
    parser.add_argument("--prof_gpu", action="store_true")

    args = parser.parse_args()

    with th.autograd.profiler.profile(
        enabled=args.th_prof, use_cuda=True, record_shapes=True, profile_memory=True, with_stack=True
    ) as prof:
        with profiler.record_function("init"):
            if args.balanced_consistency > 0 or args.contrastive > 0:
                args.augment = True
            args.name = os.path.splitext(os.path.basename(args.path))[0]
            args.r1 = args.r1 * args.size ** 2

            args.num_gpus = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1
            th.backends.cudnn.benchmark = args.cudnn_benchmark
            args.distributed = args.num_gpus > 1

            if args.distributed:
                th.cuda.set_device(args.local_rank)
                th.distributed.init_process_group(backend="nccl", init_method="env://")
                synchronize()

            generator = Generator(
                args.size,
                args.latent_size,
                args.n_mlp,
                channel_multiplier=args.channel_multiplier,
                constant_input=args.constant_input,
                min_rgb_size=args.min_rgb_size,
            ).to(device, non_blocking=True)
            discriminator = Discriminator(
                args.size, channel_multiplier=args.channel_multiplier, use_skip=args.d_skip
            ).to(device, non_blocking=True)

            if args.log_spec_norm:
                for name, parameter in generator.named_parameters():
                    if "weight" in name and parameter.squeeze().dim() > 1:
                        mod = generator
                        for attr in name.replace(".weight", "").split("."):
                            mod = getattr(mod, attr)
                        validation.track_spectral_norm(mod)
                for name, parameter in discriminator.named_parameters():
                    if "weight" in name and parameter.squeeze().dim() > 1:
                        mod = discriminator
                        for attr in name.replace(".weight", "").split("."):
                            mod = getattr(mod, attr)
                        validation.track_spectral_norm(mod)

            g_ema = Generator(
                args.size,
                args.latent_size,
                args.n_mlp,
                channel_multiplier=args.channel_multiplier,
                constant_input=args.constant_input,
                min_rgb_size=args.min_rgb_size,
            ).to(device, non_blocking=True)
            g_ema.requires_grad_(False)
            g_ema.eval()
            accumulate(g_ema, generator, 0)

            if args.contrastive > 0:
                contrast_learner = ContrastiveLearner(
                    discriminator,
                    args.size,
                    augment_fn=nn.Sequential(
                        nn.ReflectionPad2d(int((math.sqrt(2) - 1) * args.size / 4)),  # zoom out
                        augs.RandomHorizontalFlip(),
                        RandomApply(augs.RandomAffine(degrees=0, translate=(0.25, 0.25), shear=(15, 15)), p=0.1),
                        RandomApply(augs.RandomRotation(180), p=0.1),
                        augs.RandomResizedCrop(size=(args.size, args.size), scale=(1, 1), ratio=(1, 1)),
                        RandomApply(
                            augs.RandomResizedCrop(size=(args.size, args.size), scale=(0.5, 0.9)), p=0.1
                        ),  # zoom in
                        RandomApply(augs.RandomErasing(), p=0.1),
                    ),
                    hidden_layer=(-1, 0),
                )
            else:
                contrast_learner = None

            g_reg_ratio = args.g_reg_every / (args.g_reg_every + 1)
            d_reg_ratio = args.d_reg_every / (args.d_reg_every + 1)

            g_optim = th.optim.Adam(
                generator.parameters(), lr=args.lr * g_reg_ratio, betas=(0 ** g_reg_ratio, 0.99 ** g_reg_ratio),
            )
            d_optim = th.optim.Adam(
                discriminator.parameters(),
                lr=args.lr * d_reg_ratio * args.d_lr_ratio,
                betas=(0 ** d_reg_ratio, 0.99 ** d_reg_ratio),
            )

            if args.lookahead:
                g_optim = LookaheadMinimax(
                    g_optim, d_optim, la_steps=args.la_steps, la_alpha=args.la_alpha, accumulate=args.num_accumulate
                )

            if args.checkpoint is not None:
                print("load model:", args.checkpoint)

                checkpoint = th.load(args.checkpoint, map_location=lambda storage, loc: storage)

                try:
                    ckpt_name = os.path.basename(args.checkpoint)
                    args.start_iter = int(os.path.splitext(ckpt_name)[-1].replace(args.name, ""))
                except ValueError:
                    pass

                if args.transfer_mapping_only:
                    print("Using generator latent mapping network from checkpoint")
                    mapping_state_dict = {}
                    for key, val in checkpoint["state_dict"].items():
                        if "generator.style" in key:
                            mapping_state_dict[key.replace("generator.", "")] = val
                    generator.load_state_dict(mapping_state_dict, strict=False)
                else:
                    generator.load_state_dict(checkpoint["g"], strict=False)
                    g_ema.load_state_dict(checkpoint["g_ema"], strict=False)

                    discriminator.load_state_dict(checkpoint["d"], strict=False)

                    if args.lookahead:
                        g_optim.load_state_dict(checkpoint["g_optim"], checkpoint["d_optim"])
                    else:
                        g_optim.load_state_dict(checkpoint["g_optim"])
                        d_optim.load_state_dict(checkpoint["d_optim"])

                del checkpoint
                th.cuda.empty_cache()

            if args.distributed:
                generator = nn.parallel.DistributedDataParallel(
                    generator,
                    device_ids=[args.local_rank],
                    output_device=args.local_rank,
                    broadcast_buffers=False,
                    find_unused_parameters=True,
                )

                discriminator = nn.parallel.DistributedDataParallel(
                    discriminator,
                    device_ids=[args.local_rank],
                    output_device=args.local_rank,
                    broadcast_buffers=False,
                    find_unused_parameters=True,
                )

                if contrast_learner is not None:
                    contrast_learner = nn.parallel.DistributedDataParallel(
                        contrast_learner,
                        device_ids=[args.local_rank],
                        output_device=args.local_rank,
                        broadcast_buffers=False,
                        find_unused_parameters=True,
                    )

            transform = transforms.Compose(
                [
                    transforms.RandomVerticalFlip(p=0.5 if args.vflip else 0),
                    transforms.RandomHorizontalFlip(p=0.5 if args.hflip else 0),
                    transforms.ToTensor(),
                    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
                ]
            )

            dataset = MultiResolutionDataset(args.path, transform, args.size)
            loader = data.DataLoader(
                dataset,
                batch_size=args.batch_size,
                sampler=data_sampler(dataset, shuffle=True, distributed=args.distributed),
                num_workers=0,
                drop_last=True,
                pin_memory=True,
            )

            if get_rank() == 0:
                validation.get_dataset_inception_features(loader, args.name, args.size)
                if args.wbgroup is None:
                    wandb.init(project=args.wbproj, name=args.wbname, config=vars(args))
                else:
                    wandb.init(project=args.wbproj, group=args.wbgroup, name=args.wbname, config=vars(args))

        if args.prof_gpu:
            os.environ["GPU_DEBUG"] = str(args.local_rank)
            import sys
            from gpu_profile import gpu_profile

            sys.settrace(gpu_profile)
        train(args, loader, generator, discriminator, contrast_learner, g_optim, d_optim, g_ema)

    if args.th_prof:
        print(prof.total_average())
        print("cuda_memory_usage", prof.table(sort_by="cuda_memory_usage", row_limit=20))
        prof.export_chrome_trace(f"{args.name}_gpu{args.local_rank}.trace")


================================================
FILE: validation/__init__.py
================================================
from .metrics import vae_fid, fid, get_dataset_inception_features, ppl, prdc
from .spectral_norm import track_spectral_norm


================================================
FILE: validation/calc_fid.py
================================================
import argparse
import pickle

import torch
from torch import nn
import numpy as np
from scipy import linalg
from tqdm import tqdm

from model import Generator
from inception import InceptionV3


@torch.no_grad()
def extract_feature_from_samples(generator, inception, truncation, truncation_latent, batch_size, n_sample, device):
    n_batch = n_sample // batch_size
    resid = n_sample - (n_batch * batch_size)
    batch_sizes = [batch_size] * n_batch + [resid]
    features = []

    for batch in tqdm(batch_sizes):
        latent = torch.randn(batch, 512, device=device)
        img, _ = g([latent], truncation=truncation, truncation_latent=truncation_latent)
        feat = inception(img)[0].view(img.shape[0], -1)
        features.append(feat.to("cpu"))

    features = torch.cat(features, 0)

    return features


def calc_fid(sample_mean, sample_cov, real_mean, real_cov, eps=1e-6):
    cov_sqrt, _ = linalg.sqrtm(sample_cov @ real_cov, disp=False)

    if not np.isfinite(cov_sqrt).all():
        print("product of cov matrices is singular")
        offset = np.eye(sample_cov.shape[0]) * eps
        cov_sqrt = linalg.sqrtm((sample_cov + offset) @ (real_cov + offset))

    if np.iscomplexobj(cov_sqrt):
        if not np.allclose(np.diagonal(cov_sqrt).imag, 0, atol=1e-3):
            m = np.max(np.abs(cov_sqrt.imag))

            raise ValueError(f"Imaginary component {m}")

        cov_sqrt = cov_sqrt.real

    mean_diff = sample_mean - real_mean
    mean_norm = mean_diff @ mean_diff

    trace = np.trace(sample_cov) + np.trace(real_cov) - 2 * np.trace(cov_sqrt)

    fid = mean_norm + trace

    return fid


if __name__ == "__main__":
    device = "cuda"

    parser = argparse.ArgumentParser()

    parser.add_argument("--truncation", type=float, default=1)
    parser.add_argument("--truncation_mean", type=int, default=4096 * 8)
    parser.add_argument("--batch", type=int, default=64)
    parser.add_argument("--n_sample", type=int, default=50000)
    parser.add_argument("--size", type=int, default=256)
    parser.add_argument("--inception", type=str, default=None, required=True)
    parser.add_argument("ckpt", metavar="CHECKPOINT")

    args = parser.parse_args()

    ckpt = torch.load(args.ckpt)

    g = Generator(args.size, 512, 8).to(device)
    g.load_state_dict(ckpt["g_ema"])
    g = nn.DataParallel(g)
    g.eval()

    if args.truncation < 1:
        with torch.no_grad():
            mean_latent = g.mean_latent(args.truncation_mean)

    else:
        mean_latent = None

    inception = InceptionV3([3], normalize_input=False, init_weights=False)
    inception = nn.DataParallel(inception).eval().cuda()

    features = extract_feature_from_samples(
        g, inception, args.truncation, mean_latent, args.batch, args.n_sample, device
    ).numpy()
    print(f"extracted {features.shape[0]} features")

    sample_mean = np.mean(features, 0)
    sample_cov = np.cov(features, rowvar=False)

    with open(args.inception, "rb") as f:
        embeds = pickle.load(f)
        real_mean = embeds["mean"]
        real_cov = embeds["cov"]

    fid = calc_fid(sample_mean, sample_cov, real_mean, real_cov)

    print("fid:", fid)


================================================
FILE: validation/calc_inception.py
================================================
import argparse
import pickle
import os

import torch
from torch import nn
from torch.nn import functional as F
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.models import Inception3
import numpy as np
from tqdm import tqdm

from inception import InceptionV3
from dataset import MultiResolutionDataset


class Inception3Feature(Inception3):
    def forward(self, x):
        if x.shape[2] != 299 or x.shape[3] != 299:
            x = F.interpolate(x, size=(299, 299), mode="bilinear", align_corners=True)

        x = self.Conv2d_1a_3x3(x)  # 299 x 299 x 3
        x = self.Conv2d_2a_3x3(x)  # 149 x 149 x 32
        x = self.Conv2d_2b_3x3(x)  # 147 x 147 x 32
        x = F.max_pool2d(x, kernel_size=3, stride=2)  # 147 x 147 x 64

        x = self.Conv2d_3b_1x1(x)  # 73 x 73 x 64
        x = self.Conv2d_4a_3x3(x)  # 73 x 73 x 80
        x = F.max_pool2d(x, kernel_size=3, stride=2)  # 71 x 71 x 192

        x = self.Mixed_5b(x)  # 35 x 35 x 192
        x = self.Mixed_5c(x)  # 35 x 35 x 256
        x = self.Mixed_5d(x)  # 35 x 35 x 288

        x = self.Mixed_6a(x)  # 35 x 35 x 288
        x = self.Mixed_6b(x)  # 17 x 17 x 768
        x = self.Mixed_6c(x)  # 17 x 17 x 768
        x = self.Mixed_6d(x)  # 17 x 17 x 768
        x = self.Mixed_6e(x)  # 17 x 17 x 768

        x = self.Mixed_7a(x)  # 17 x 17 x 768
        x = self.Mixed_7b(x)  # 8 x 8 x 1280
        x = self.Mixed_7c(x)  # 8 x 8 x 2048

        x = F.avg_pool2d(x, kernel_size=8)  # 8 x 8 x 2048

        return x.view(x.shape[0], x.shape[1])  # 1 x 1 x 2048


def load_patched_inception_v3():
    # inception = inception_v3(pretrained=True)
    # inception_feat = Inception3Feature()
    # inception_feat.load_state_dict(inception.state_dict())
    inception_feat = InceptionV3([3], normalize_input=False, init_weights=False)

    return inception_feat


@torch.no_grad()
def extract_features(loader, inception, device):
    pbar = tqdm(loader)

    feature_list = []

    for img in pbar:
        img = img.to(device)
        feature = inception(img)[0].view(img.shape[0], -1)
        feature_list.append(feature.to("cpu"))

    features = torch.cat(feature_list, 0)

    return features


if __name__ == "__main__":
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    parser = argparse.ArgumentParser(description="Calculate Inception v3 features for datasets")
    parser.add_argument("--size", type=int, default=256)
    parser.add_argument("--batch", default=64, type=int, help="batch size")
    parser.add_argument("--n_sample", type=int, default=50000)
    parser.add_argument("--vflip", action="store_true")
    parser.add_argument("--hflip", action="store_true")
    parser.add_argument("path", metavar="PATH", help="path to datset lmdb file")

    args = parser.parse_args()

    inception = load_patched_inception_v3()
    inception = nn.DataParallel(inception).eval().to(device)

    transform = transforms.Compose(
        [
            transforms.RandomVerticalFlip(p=0.5 if args.vflip else 0),
            transforms.RandomHorizontalFlip(p=0.5 if args.hflip else 0),
            transforms.ToTensor(),
            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
        ]
    )

    dset = MultiResolutionDataset(args.path, transform=transform, resolution=args.size)
    loader = DataLoader(dset, batch_size=args.batch, num_workers=4)

    features = extract_features(loader, inception, device).numpy()

    features = features[: args.n_sample]

    print(f"extracted {features.shape[0]} features")

    mean = np.mean(features, 0)
    cov = np.cov(features, rowvar=False)

    name = os.path.splitext(os.path.basename(args.path))[0]

    with open(f"inception_{name}.pkl", "wb") as f:
        pickle.dump({"mean": mean, "cov": cov, "size": args.size, "path": args.path}, f)


================================================
FILE: validation/calc_ppl.py
================================================
import argparse

import torch
from torch.nn import functional as F
import numpy as np
from tqdm import tqdm

import lpips
from model import Generator


def normalize(x):
    return x / torch.sqrt(x.pow(2).sum(-1, keepdim=True))


def slerp(a, b, t):
    a = normalize(a)
    b = normalize(b)
    d = (a * b).sum(-1, keepdim=True)
    p = t * torch.acos(d)
    c = normalize(b - d * a)
    d = a * torch.cos(p) + c * torch.sin(p)

    return normalize(d)


def lerp(a, b, t):
    return a + (b - a) * t


if __name__ == "__main__":
    device = "cuda"

    parser = argparse.ArgumentParser()

    parser.add_argument("--space", choices=["z", "w"])
    parser.add_argument("--batch", type=int, default=64)
    parser.add_argument("--n_sample", type=int, default=5000)
    parser.add_argument("--size", type=int, default=256)
    parser.add_argument("--eps", type=float, default=1e-4)
    parser.add_argument("--crop", action="store_true")
    parser.add_argument("ckpt", metavar="CHECKPOINT")

    args = parser.parse_args()

    latent_dim = 512

    ckpt = torch.load(args.ckpt)

    g = Generator(args.size, latent_dim, 8).to(device)
    g.load_state_dict(ckpt["g_ema"])
    g.eval()

    percept = lpips.PerceptualLoss(model="net-lin", net="vgg", use_gpu=device.startswith("cuda"))

    distances = []

    n_batch = args.n_sample // args.batch
    resid = args.n_sample - (n_batch * args.batch)
    batch_sizes = [args.batch] * n_batch + [resid]

    with torch.no_grad():
        for batch in tqdm(batch_sizes):
            noise = g.make_noise()

            inputs = torch.randn([batch * 2, latent_dim], device=device)
            lerp_t = torch.rand(batch, device=device)

            if args.space == "w":
                latent = g.get_latent(inputs)
                latent_t0, latent_t1 = latent[::2], latent[1::2]
                latent_e0 = lerp(latent_t0, latent_t1, lerp_t[:, None])
                latent_e1 = lerp(latent_t0, latent_t1, lerp_t[:, None] + args.eps)
                latent_e = torch.stack([latent_e0, latent_e1], 1).view(*latent.shape)

            image, _ = g([latent_e], input_is_latent=True, noise=noise)

            if args.crop:
                c = image.shape[2] // 8
                image = image[:, :, c * 3 : c * 7, c * 2 : c * 6]

            factor = image.shape[2] // 256

            if factor > 1:
                image = F.interpolate(image, size=(256, 256), mode="bilinear", align_corners=False)

            dist = percept(image[::2], image[1::2]).view(image.shape[0] // 2) / (args.eps ** 2)
            distances.append(dist.to("cpu").numpy())

    distances = np.concatenate(distances, 0)

    lo = np.percentile(distances, 1, interpolation="lower")
    hi = np.percentile(distances, 99, interpolation="higher")
    filtered_dist = np.extract(np.logical_and(lo <= distances, distances <= hi), distances)

    print("ppl:", filtered_dist.mean())


================================================
FILE: validation/inception.py
================================================
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import models

try:
    from torchvision.models.utils import load_state_dict_from_url
except ImportError:
    from torch.utils.model_zoo import load_url as load_state_dict_from_url

# Inception weights ported to Pytorch from
# http://download.tensorflow.org/models/image/imagenet/inception-2015-12-05.tgz
FID_WEIGHTS_URL = (
    "https://github.com/mseitzer/pytorch-fid/releases/download/fid_weights/pt_inception-2015-12-05-6726825d.pth"
)


class InceptionV3(nn.Module):
    """Pretrained InceptionV3 network returning feature maps"""

    # Index of default block of inception to return,
    # corresponds to output of final average pooling
    DEFAULT_BLOCK_INDEX = 3

    # Maps feature dimensionality to their output blocks indices
    BLOCK_INDEX_BY_DIM = {
        64: 0,  # First max pooling features
        192: 1,  # Second max pooling featurs
        768: 2,  # Pre-aux classifier features
        2048: 3,  # Final average pooling features
    }

    def __init__(
        self,
        output_blocks=[DEFAULT_BLOCK_INDEX],
        resize_input=True,
        normalize_input=True,
        requires_grad=False,
        init_weights=False,
        use_fid_inception=False,
    ):
        """Build pretrained InceptionV3

        Parameters
        ----------
        output_blocks : list of int
            Indices of blocks to return features of. Possible values are:
                - 0: corresponds to output of first max pooling
                - 1: corresponds to output of second max pooling
                - 2: corresponds to output which is fed to aux classifier
                - 3: corresponds to output of final average pooling
        resize_input : bool
            If true, bilinearly resizes input to width and height 299 before
            feeding input to model. As the network without fully connected
            layers is fully convolutional, it should be able to handle inputs
            of arbitrary size, so resizing might not be strictly needed
        normalize_input : bool
            If true, scales the input from range (0, 1) to the range the
            pretrained Inception network expects, namely (-1, 1)
        requires_grad : bool
            If true, parameters of the model require gradients. Possibly useful
            for finetuning the network
        use_fid_inception : bool
            If true, uses the pretrained Inception model used in Tensorflow's
            FID implementation. If false, uses the pretrained Inception model
            available in torchvision. The FID Inception model has different
            weights and a slightly different structure from torchvision's
            Inception model. If you want to compute FID scores, you are
            strongly advised to set this parameter to true to get comparable
            results.
        """
        super(InceptionV3, self).__init__()

        self.resize_input = resize_input
        self.normalize_input = normalize_input
        self.output_blocks = sorted(output_blocks)
        self.last_needed_block = max(output_blocks)

        assert self.last_needed_block <= 3, "Last possible output block index is 3"

        self.blocks = nn.ModuleList()

        if use_fid_inception:
            inception = fid_inception_v3()
        else:
            inception = models.inception_v3(pretrained=True)  # , init_weights=False)

        # Block 0: input to maxpool1
        block0 = [
            inception.Conv2d_1a_3x3,
            inception.Conv2d_2a_3x3,
            inception.Conv2d_2b_3x3,
            nn.MaxPool2d(kernel_size=3, stride=2),
        ]
        self.blocks.append(nn.Sequential(*block0))

        # Block 1: maxpool1 to maxpool2
        if self.last_needed_block >= 1:
            block1 = [inception.Conv2d_3b_1x1, inception.Conv2d_4a_3x3, nn.MaxPool2d(kernel_size=3, stride=2)]
            self.blocks.append(nn.Sequential(*block1))

        # Block 2: maxpool2 to aux classifier
        if self.last_needed_block >= 2:
            block2 = [
                inception.Mixed_5b,
                inception.Mixed_5c,
                inception.Mixed_5d,
                inception.Mixed_6a,
                inception.Mixed_6b,
                inception.Mixed_6c,
                inception.Mixed_6d,
                inception.Mixed_6e,
            ]
            self.blocks.append(nn.Sequential(*block2))

        # Block 3: aux classifier to final avgpool
        if self.last_needed_block >= 3:
            block3 = [
                inception.Mixed_7a,
                inception.Mixed_7b,
                inception.Mixed_7c,
                nn.AdaptiveAvgPool2d(output_size=(1, 1)),
            ]
            self.blocks.append(nn.Sequential(*block3))

        for param in self.parameters():
            param.requires_grad = requires_grad

    def forward(self, inp):
        """Get Inception feature maps

        Parameters
        ----------
        inp : torch.autograd.Variable
            Input tensor of shape Bx3xHxW. Values are expected to be in
            range (0, 1)

        Returns
        -------
        List of torch.autograd.Variable, corresponding to the selected output
        block, sorted ascending by index
        """
        outp = []
        x = inp

        if self.resize_input:
            x = F.interpolate(x, size=(299, 299), mode="bilinear", align_corners=False)

        if self.normalize_input:
            x = 2 * x - 1  # Scale from range (0, 1) to range (-1, 1)

        for idx, block in enumerate(self.blocks):
            x = block(x)
            if idx in self.output_blocks:
                outp.append(x)

            if idx == self.last_needed_block:
                break

        return outp


def fid_inception_v3():
    """Build pretrained Inception model for FID computation

    The Inception model for FID computation uses a different set of weights
    and has a slightly different structure than torchvision's Inception.

    This method first constructs torchvision's Inception and then patches the
    necessary parts that are different in the FID Inception model.
    """
    inception = models.inception_v3(num_classes=1008, aux_logits=False, pretrained=False)
    inception.Mixed_5b = FIDInceptionA(192, pool_features=32)
    inception.Mixed_5c = FIDInceptionA(256, pool_features=64)
    inception.Mixed_5d = FIDInceptionA(288, pool_features=64)
    inception.Mixed_6b = FIDInceptionC(768, channels_7x7=128)
    inception.Mixed_6c = FIDInceptionC(768, channels_7x7=160)
    inception.Mixed_6d = FIDInceptionC(768, channels_7x7=160)
    inception.Mixed_6e = FIDInceptionC(768, channels_7x7=192)
    inception.Mixed_7b = FIDInceptionE_1(1280)
    inception.Mixed_7c = FIDInceptionE_2(2048)

    state_dict = load_state_dict_from_url(FID_WEIGHTS_URL, progress=True)
    inception.load_state_dict(state_dict)
    return inception


class FIDInceptionA(models.inception.InceptionA):
    """InceptionA block patched for FID computation"""

    def __init__(self, in_channels, pool_features):
        super(FIDInceptionA, self).__init__(in_channels, pool_features)

    def forward(self, x):
        branch1x1 = self.branch1x1(x)

        branch5x5 = self.branch5x5_1(x)
        branch5x5 = self.branch5x5_2(branch5x5)

        branch3x3dbl = self.branch3x3dbl_1(x)
        branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)
        branch3x3dbl = self.branch3x3dbl_3(branch3x3dbl)

        # Patch: Tensorflow's average pool does not use the padded zero's in
        # its average calculation
        branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1, count_include_pad=False)
        branch_pool = self.branch_pool(branch_pool)

        outputs = [branch1x1, branch5x5, branch3x3dbl, branch_pool]
        return torch.cat(outputs, 1)


class FIDInceptionC(models.inception.InceptionC):
    """InceptionC block patched for FID computation"""

    def __init__(self, in_channels, channels_7x7):
        super(FIDInceptionC, self).__init__(in_channels, channels_7x7)

    def forward(self, x):
        branch1x1 = self.branch1x1(x)

        branch7x7 = self.branch7x7_1(x)
        branch7x7 = self.branch7x7_2(branch7x7)
        branch7x7 = self.branch7x7_3(branch7x7)

        branch7x7dbl = self.branch7x7dbl_1(x)
        branch7x7dbl = self.branch7x7dbl_2(branch7x7dbl)
        branch7x7dbl = self.branch7x7dbl_3(branch7x7dbl)
        branch7x7dbl = self.branch7x7dbl_4(branch7x7dbl)
        branch7x7dbl = self.branch7x7dbl_5(branch7x7dbl)

        # Patch: Tensorflow's average pool does not use the padded zero's in
        # its average calculation
        branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1, count_include_pad=False)
        branch_pool = self.branch_pool(branch_pool)

        outputs = [branch1x1, branch7x7, branch7x7dbl, branch_pool]
        return torch.cat(outputs, 1)


class FIDInceptionE_1(models.inception.InceptionE):
    """First InceptionE block patched for FID computation"""

    def __init__(self, in_channels):
        super(FIDInceptionE_1, self).__init__(in_channels)

    def forward(self, x):
        branch1x1 = self.branch1x1(x)

        branch3x3 = self.branch3x3_1(x)
        branch3x3 = [
            self.branch3x3_2a(branch3x3),
            self.branch3x3_2b(branch3x3),
        ]
        branch3x3 = torch.cat(branch3x3, 1)

        branch3x3dbl = self.branch3x3dbl_1(x)
        branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)
        branch3x3dbl = [
            self.branch3x3dbl_3a(branch3x3dbl),
            self.branch3x3dbl_3b(branch3x3dbl),
        ]
        branch3x3dbl = torch.cat(branch3x3dbl, 1)

        # Patch: Tensorflow's average pool does not use the padded zero's in
        # its average calculation
        branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1, count_include_pad=False)
        branch_pool = self.branch_pool(branch_pool)

        outputs = [branch1x1, branch3x3, branch3x3dbl, branch_pool]
        return torch.cat(outputs, 1)


class FIDInceptionE_2(models.inception.InceptionE):
    """Second InceptionE block patched for FID computation"""

    def __init__(self, in_channels):
        super(FIDInceptionE_2, self).__init__(in_channels)

    def forward(self, x):
        branch1x1 = self.branch1x1(x)

        branch3x3 = self.branch3x3_1(x)
        branch3x3 = [
            self.branch3x3_2a(branch3x3),
            self.branch3x3_2b(branch3x3),
        ]
        branch3x3 = torch.cat(branch3x3, 1)

        branch3x3dbl = self.branch3x3dbl_1(x)
        branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)
        branch3x3dbl = [
            self.branch3x3dbl_3a(branch3x3dbl),
            self.branch3x3dbl_3b(branch3x3dbl),
        ]
        branch3x3dbl = torch.cat(branch3x3dbl, 1)

        # Patch: The FID Inception model uses max pooling instead of average
        # pooling. This is likely an error in this specific Inception
        # implementation, as other Inception models use average pooling here
        # (which matches the description in the paper).
        branch_pool = F.max_pool2d(x, kernel_size=3, stride=1, padding=1)
        branch_pool = self.branch_pool(branch_pool)

        outputs = [branch1x1, branch3x3, branch3x3dbl, branch_pool]
        return torch.cat(outputs, 1)


================================================
FILE: validation/lpips/__init__.py
================================================
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
from skimage.measure import compare_ssim
import torch
from torch.autograd import Variable

from . import dist_model


class PerceptualLoss(torch.nn.Module):
    def __init__(
        self, model="net-lin", net="alex", colorspace="rgb", spatial=False, use_gpu=True, gpu_ids=[0]
    ):  # VGG using our perceptually-learned weights (LPIPS metric)
        # def __init__(self, model='net', net='vgg', use_gpu=True): # "default" way of using VGG as a perceptual loss
        super(PerceptualLoss, self).__init__()
        # print('Setting up Perceptual loss...')
        self.use_gpu = use_gpu
        self.spatial = spatial
        self.gpu_ids = gpu_ids
        self.model = dist_model.DistModel()
        self.model.initialize(
            model=model, net=net, use_gpu=use_gpu, colorspace=colorspace, spatial=self.spatial, gpu_ids=gpu_ids
        )
        # print('...[%s] initialized'%self.model.name())
        # print('...Done')

    def forward(self, pred, target, normalize=False):
        """
        Pred and target are Variables.
        If normalize is True, assumes the images are between [0,1] and then scales them between [-1,+1]
        If normalize is False, assumes the images are already between [-1,+1]

        Inputs pred and target are Nx3xHxW
        Output pytorch Variable N long
        """

        if normalize:
            target = 2 * target - 1
            pred = 2 * pred - 1

        return self.model.forward(target, pred)


================================================
FILE: validation/lpips/base_model.py
================================================
import os
import torch


class BaseModel:
    def __init__(self):
        pass

    def name(self):
        return "BaseModel"

    def initialize(self, use_gpu=True, gpu_ids=[0]):
        self.use_gpu = use_gpu
        self.gpu_ids = gpu_ids

    def forward(self):
        pass

    def get_image_paths(self):
        pass

    def optimize_parameters(self):
        pass

    def get_current_visuals(self):
        return self.input

    def get_current_errors(self):
        return {}

    def save(self, label):
        pass

    # helper saving function that can be used by subclasses
    def save_network(self, network, path, network_label, epoch_label):
        save_filename = "%s_net_%s.pth" % (epoch_label, network_label)
        save_path = os.path.join(path, save_filename)
        torch.save(network.state_dict(), save_path)

    # helper loading function that can be used by subclasses
    def load_network(self, network, network_label, epoch_label):
        save_filename = "%s_net_%s.pth" % (epoch_label, network_label)
        save_path = os.path.join(self.save_dir, save_filename)
        print("Loading network from %s" % save_path)
        network.load_state_dict(torch.load(save_path))

    def update_learning_rate():
        pass

    def get_image_paths(self):
        return self.image_paths

    def save_done(self, flag=False):
        np.save(os.path.join(self.save_dir, "done_flag"), flag)
        np.savetxt(os.path.join(self.save_dir, "done_flag"), [flag,], fmt="%i")



================================================
FILE: validation/lpips/dist_model.py
================================================
import numpy as np
import torch
import os
from collections import OrderedDict
from torch.autograd import Variable
from .base_model import BaseModel
from scipy.ndimage import zoom
from tqdm import tqdm


from . import networks_basic as networks
from . import util


class DistModel(BaseModel):
    def name(self):
        return self.model_name

    def initialize(
        self,
        model="net-lin",
        net="alex",
        colorspace="Lab",
        pnet_rand=False,
        pnet_tune=False,
        model_path=None,
        use_gpu=True,
        printNet=False,
        spatial=False,
        is_train=False,
        lr=0.0001,
        beta1=0.5,
        version="0.1",
        gpu_ids=[0],
    ):
        """
        INPUTS
            model - ['net-lin'] for linearly calibrated network
                    ['net'] for off-the-shelf network
                    ['L2'] for L2 distance in Lab colorspace
                    ['SSIM'] for ssim in RGB colorspace
            net - ['squeeze','alex','vgg']
            model_path - if None, will look in weights/[NET_NAME].pth
            colorspace - ['Lab','RGB'] colorspace to use for L2 and SSIM
            use_gpu - bool - whether or not to use a GPU
            printNet - bool - whether or not to print network architecture out
            spatial - bool - whether to output an array containing varying distances across spatial dimensions
            spatial_shape - if given, output spatial shape. if None then spatial shape is determined automatically via spatial_factor (see below).
            spatial_factor - if given, specifies upsampling factor relative to the largest spatial extent of a convolutional layer. if None then resized to size of input images.
            spatial_order - spline order of filter for upsampling in spatial mode, by default 1 (bilinear).
            is_train - bool - [True] for training mode
            lr - float - initial learning rate
            beta1 - float - initial momentum term for adam
            version - 0.1 for latest, 0.0 was original (with a bug)
            gpu_ids - int array - [0] by default, gpus to use
        """
        BaseModel.initialize(self, use_gpu=use_gpu, gpu_ids=gpu_ids)

        self.model = model
        self.net = net
        self.is_train = is_train
        self.spatial = spatial
        self.gpu_ids = gpu_ids
        self.model_name = "%s [%s]" % (model, net)

        if self.model == "net-lin":  # pretrained net + linear layer
            self.net = networks.PNetLin(
                pnet_rand=pnet_rand,
                pnet_tune=pnet_tune,
                pnet_type=net,
                use_dropout=True,
                spatial=spatial,
                version=version,
                lpips=True,
            )
            kw = {}
            if not use_gpu:
                kw["map_location"] = "cpu"
            if model_path is None:
                import inspect

                model_path = os.path.abspath(
                    os.path.join(inspect.getfile(self.initialize), "..", "weights/v%s/%s.pth" % (version, net))
                )

            if not is_train:
                # print("Loading model from: %s" % model_path)
                self.net.load_state_dict(torch.load(model_path, **kw), strict=False)

        elif self.model == "net":  # pretrained network
            self.net = networks.PNetLin(pnet_rand=pnet_rand, pnet_type=net, lpips=False)
        elif self.model in ["L2", "l2"]:
            self.net = networks.L2(use_gpu=use_gpu, colorspace=colorspace)  # not really a network, only for testing
            self.model_name = "L2"
        elif self.model in ["DSSIM", "dssim", "SSIM", "ssim"]:
            self.net = networks.DSSIM(use_gpu=use_gpu, colorspace=colorspace)
            self.model_name = "SSIM"
        else:
            raise ValueError("Model [%s] not recognized." % self.model)

        self.parameters = list(self.net.parameters())

        if self.is_train:  # training mode
            # extra network on top to go from distances (d0,d1) => predicted human judgment (h*)
            self.rankLoss = networks.BCERankingLoss()
            self.parameters += list(self.rankLoss.net.parameters())
            self.lr = lr
            self.old_lr = lr
            self.optimizer_net = torch.optim.Adam(self.parameters, lr=lr, betas=(beta1, 0.999))
        else:  # test mode
            self.net.eval()

        if use_gpu:
            self.net.to(gpu_ids[0])
            self.net = torch.nn.DataParallel(self.net, device_ids=gpu_ids)
            if self.is_train:
                self.rankLoss = self.rankLoss.to(device=gpu_ids[0])  # just put this on GPU0

        if printNet:
            print("---------- Networks initialized -------------")
            networks.print_network(self.net)
            print("-----------------------------------------------")

    def forward(self, in0, in1, retPerLayer=False):
        """ Function computes the distance between image patches in0 and in1
        INPUTS
            in0, in1 - torch.Tensor object of shape Nx3xXxY - image patch scaled to [-1,1]
        OUTPUT
            computed distances between in0 and in1
        """

        return self.net.forward(in0, in1, retPerLayer=retPerLayer)

    # ***** TRAINING FUNCTIONS *****
    def optimize_parameters(self):
        self.forward_train()
        self.optimizer_net.zero_grad()
        self.backward_train()
        self.optimizer_net.step()
        self.clamp_weights()

    def clamp_weights(self):
        for module in self.net.modules():
            if hasattr(module, "weight") and module.kernel_size == (1, 1):
                module.weight.data = torch.clamp(module.weight.data, min=0)

    def set_input(self, data):
        self.input_ref = data["ref"]
        self.input_p0 = data["p0"]
        self.input_p1 = data["p1"]
        self.input_judge = data["judge"]

        if self.use_gpu:
            self.input_ref = self.input_ref.to(device=self.gpu_ids[0])
            self.input_p0 = self.input_p0.to(device=self.gpu_ids[0])
            self.input_p1 = self.input_p1.to(device=self.gpu_ids[0])
            self.input_judge = self.input_judge.to(device=self.gpu_ids[0])

        self.var_ref = Variable(self.input_ref, requires_grad=True)
        self.var_p0 = Variable(self.input_p0, requires_grad=True)
        self.var_p1 = Variable(self.input_p1, requires_grad=True)

    def forward_train(self):  # run forward pass
        # print(self.net.module.scaling_layer.shift)
        # print(torch.norm(self.net.module.net.slice1[0].weight).item(), torch.norm(self.net.module.lin0.model[1].weight).item())

        self.d0 = self.forward(self.var_ref, self.var_p0)
        self.d1 = self.forward(self.var_ref, self.var_p1)
        self.acc_r = self.compute_accuracy(self.d0, self.d1, self.input_judge)

        self.var_judge = Variable(1.0 * self.input_judge).view(self.d0.size())

        self.loss_total = self.rankLoss.forward(self.d0, self.d1, self.var_judge * 2.0 - 1.0)

        return self.loss_total

    def backward_train(self):
        torch.mean(self.loss_total).backward()

    def compute_accuracy(self, d0, d1, judge):
        """ d0, d1 are Variables, judge is a Tensor """
        d1_lt_d0 = (d1 < d0).cpu().data.numpy().flatten()
        judge_per = judge.cpu().numpy().flatten()
        return d1_lt_d0 * judge_per + (1 - d1_lt_d0) * (1 - judge_per)

    def get_current_errors(self):
        retDict = OrderedDict([("loss_total", self.loss_total.data.cpu().numpy()), ("acc_r", self.acc_r)])

        for key in retDict.keys():
            retDict[key] = np.mean(retDict[key])

        return retDict

    def get_current_visuals(self):
        zoom_factor = 256 / self.var_ref.data.size()[2]

        ref_img = util.tensor2im(self.var_ref.data)
        p0_img = util.tensor2im(self.var_p0.data)
        p1_img = util.tensor2im(self.var_p1.data)

        ref_img_vis = zoom(ref_img, [zoom_factor, zoom_factor, 1], order=0)
        p0_img_vis = zoom(p0_img, [zoom_factor, zoom_factor, 1], order=0)
        p1_img_vis = zoom(p1_img, [zoom_factor, zoom_factor, 1], order=0)

        return OrderedDict([("ref", ref_img_vis), ("p0", p0_img_vis), ("p1", p1_img_vis)])

    def save(self, path, label):
        if self.use_gpu:
            self.save_network(self.net.module, path, "", label)
        else:
            self.save_network(self.net, path, "", label)
        self.save_network(self.rankLoss.net, path, "rank", label)

    def update_learning_rate(self, nepoch_decay):
        lrd = self.lr / nepoch_decay
        lr = self.old_lr - lrd

        for param_group in self.optimizer_net.param_groups:
            param_group["lr"] = lr

        print("update lr [%s] decay: %f -> %f" % (type, self.old_lr, lr))
        self.old_lr = lr


def score_2afc_dataset(data_loader, func, name=""):
    """ Function computes Two Alternative Forced Choice (2AFC) score using
        distance function 'func' in dataset 'data_loader'
    INPUTS
        data_loader - CustomDatasetDataLoader object - contains a TwoAFCDataset inside
        func - callable distance function - calling d=func(in0,in1) should take 2
            pytorch tensors with shape Nx3xXxY, and return numpy array of length N
    OUTPUTS
        [0] - 2AFC score in [0,1], fraction of time func agrees with human evaluators
        [1] - dictionary with following elements
            d0s,d1s - N arrays containing distances between reference patch to perturbed patches 
            gts - N array in [0,1], preferred patch selected by human evaluators
                (closer to "0" for left patch p0, "1" for right patch p1,
                "0.6" means 60pct people preferred right patch, 40pct preferred left)
            scores - N array in [0,1], corresponding to what percentage function agreed with humans
    CONSTS
        N - number of test triplets in data_loader
    """

    d0s = []
    d1s = []
    gts = []

    for data in tqdm(data_loader.load_data(), desc=name):
        d0s += func(data["ref"], data["p0"]).data.cpu().numpy().flatten().tolist()
        d1s += func(data["ref"], data["p1"]).data.cpu().numpy().flatten().tolist()
        gts += data["judge"].cpu().numpy().flatten().tolist()

    d0s = np.array(d0s)
    d1s = np.array(d1s)
    gts = np.array(gts)
    scores = (d0s < d1s) * (1.0 - gts) + (d1s < d0s) * gts + (d1s == d0s) * 0.5

    return (np.mean(scores), dict(d0s=d0s, d1s=d1s, gts=gts, scores=scores))


def score_jnd_dataset(data_loader, func, name=""):
    """ Function computes JND score using distance function 'func' in dataset 'data_loader'
    INPUTS
        data_loader - CustomDatasetDataLoader object - contains a JNDDataset inside
        func - callable distance function - calling d=func(in0,in1) should take 2
            pytorch tensors with shape Nx3xXxY, and return pytorch array of length N
    OUTPUTS
        [0] - JND score in [0,1], mAP score (area under precision-recall curve)
        [1] - dictionary with following elements
            ds - N array containing distances between two patches shown to human evaluator
            sames - N array containing fraction of people who thought the two patches were identical
    CONSTS
        N - number of test triplets in data_loader
    """

    ds = []
    gts = []

    for data in tqdm(data_loader.load_data(), desc=name):
        ds += func(data["p0"], data["p1"]).data.cpu().numpy().tolist()
        gts += data["same"].cpu().numpy().flatten().tolist()

    sames = np.array(gts)
    ds = np.array(ds)

    sorted_inds = np.argsort(ds)
    ds_sorted = ds[sorted_inds]
    sames_sorted = sames[sorted_inds]

    TPs = np.cumsum(sames_sorted)
    FPs = np.cumsum(1 - sames_sorted)
    FNs = np.sum(sames_sorted) - TPs

    precs = TPs / (TPs + FPs)
    recs = TPs / (TPs + FNs)
    score = util.voc_ap(recs, precs)

    return (score, dict(ds=ds, sames=sames))



================================================
FILE: validation/lpips/networks_basic.py
================================================
import torch
import torch.nn as nn
from torch.autograd import Variable
from . import pretrained_networks as pn

from . import util


def spatial_average(in_tens, keepdim=True):
    return in_tens.mean([2, 3], keepdim=keepdim)


def upsample(in_tens, out_H=64):  # assumes scale factor is same for H and W
    in_H = in_tens.shape[2]
    scale_factor = 1.0 * out_H / in_H

    return nn.Upsample(scale_factor=scale_factor, mode="bilinear", align_corners=False)(in_tens)


# Learned perceptual metric
class PNetLin(nn.Module):
    def __init__(
        self,
        pnet_type="vgg",
        pnet_rand=False,
        pnet_tune=False,
        use_dropout=True,
        spatial=False,
        version="0.1",
        lpips=True,
    ):
        super(PNetLin, self).__init__()

        self.pnet_type = pnet_type
        self.pnet_tune = pnet_tune
        self.pnet_rand = pnet_rand
        self.spatial = spatial
        self.lpips = lpips
        self.version = version
        self.scaling_layer = ScalingLayer()

        if self.pnet_type in ["vgg", "vgg16"]:
            net_type = pn.vgg16
            self.chns = [64, 128, 256, 512, 512]
        elif self.pnet_type == "alex":
            net_type = pn.alexnet
            self.chns = [64, 192, 384, 256, 256]
        elif self.pnet_type == "squeeze":
            net_type = pn.squeezenet
            self.chns = [64, 128, 256, 384, 384, 512, 512]
        self.L = len(self.chns)

        self.net = net_type(pretrained=not self.pnet_rand, requires_grad=self.pnet_tune)

        if lpips:
            self.lin0 = NetLinLayer(self.chns[0], use_dropout=use_dropout)
            self.lin1 = NetLinLayer(self.chns[1], use_dropout=use_dropout)
            self.lin2 = NetLinLayer(self.chns[2], use_dropout=use_dropout)
            self.lin3 = NetLinLayer(self.chns[3], use_dropout=use_dropout)
            self.lin4 = NetLinLayer(self.chns[4], use_dropout=use_dropout)
            self.lins = [self.lin0, self.lin1, self.lin2, self.lin3, self.lin4]
            if self.pnet_type == "squeeze":  # 7 layers for squeezenet
                self.lin5 = NetLinLayer(self.chns[5], use_dropout=use_dropout)
                self.lin6 = NetLinLayer(self.chns[6], use_dropout=use_dropout)
                self.lins += [self.lin5, self.lin6]

    def forward(self, in0, in1, retPerLayer=False):
        # v0.0 - original release had a bug, where input was not scaled
        in0_input, in1_input = (
            (self.scaling_layer(in0), self.scaling_layer(in1)) if self.version == "0.1" else (in0, in1)
        )
        outs0, outs1 = self.net.forward(in0_input), self.net.forward(in1_input)
        feats0, feats1, diffs = {}, {}, {}

        for kk in range(self.L):
            feats0[kk], feats1[kk] = util.normalize_tensor(outs0[kk]), util.normalize_tensor(outs1[kk])
            diffs[kk] = (feats0[kk] - feats1[kk]) ** 2

        if self.lpips:
            if self.spatial:
                res = [upsample(self.lins[kk].model(diffs[kk]), out_H=in0.shape[2]) for kk in range(self.L)]
            else:
                res = [spatial_average(self.lins[kk].model(diffs[kk]), keepdim=True) for kk in range(self.L)]
        else:
            if self.spatial:
                res = [upsample(diffs[kk].sum(dim=1, keepdim=True), out_H=in0.shape[2]) for kk in range(self.L)]
            else:
                res = [spatial_average(diffs[kk].sum(dim=1, keepdim=True), keepdim=True) for kk in range(self.L)]

        val = res[0]
        for l in range(1, self.L):
            val += res[l]

        if retPerLayer:
            return (val, res)
        else:
            return val


class ScalingLayer(nn.Module):
    def __init__(self):
        super(ScalingLayer, self).__init__()
        self.register_buffer("shift", torch.Tensor([-0.030, -0.088, -0.188])[None, :, None, None])
        self.register_buffer("scale", torch.Tensor([0.458, 0.448, 0.450])[None, :, None, None])

    def forward(self, inp):
        return (inp - self.shift) / self.scale


class NetLinLayer(nn.Module):
    """ A single linear layer which does a 1x1 conv """

    def __init__(self, chn_in, chn_out=1, use_dropout=False):
        super(NetLinLayer, self).__init__()

        layers = [nn.Dropout(),] if (use_dropout) else []
        layers += [
            nn.Conv2d(chn_in, chn_out, 1, stride=1, padding=0, bias=False),
        ]
        self.model = nn.Sequential(*layers)


class Dist2LogitLayer(nn.Module):
    """ takes 2 distances, puts through fc layers, spits out value between [0,1] (if use_sigmoid is True) """

    def __init__(self, chn_mid=32, use_sigmoid=True):
        super(Dist2LogitLayer, self).__init__()

        layers = [
            nn.Conv2d(5, chn_mid, 1, stride=1, padding=0, bias=True),
        ]
        layers += [
            nn.LeakyReLU(0.2, True),
        ]
        layers += [
            nn.Conv2d(chn_mid, chn_mid, 1, stride=1, padding=0, bias=True),
        ]
        layers += [
            nn.LeakyReLU(0.2, True),
        ]
        layers += [
            nn.Conv2d(chn_mid, 1, 1, stride=1, padding=0, bias=True),
        ]
        if use_sigmoid:
            layers += [
                nn.Sigmoid(),
            ]
        self.model = nn.Sequential(*layers)

    def forward(self, d0, d1, eps=0.1):
        return self.model.forward(torch.cat((d0, d1, d0 - d1, d0 / (d1 + eps), d1 / (d0 + eps)), dim=1))


class BCERankingLoss(nn.Module):
    def __init__(self, chn_mid=32):
        super(BCERankingLoss, self).__init__()
        self.net = Dist2LogitLayer(chn_mid=chn_mid)
        # self.parameters = list(self.net.parameters())
        self.loss = torch.nn.BCELoss()

    def forward(self, d0, d1, judge):
        per = (judge + 1.0) / 2.0
        self.logit = self.net.forward(d0, d1)
        return self.loss(self.logit, per)


# L2, DSSIM metrics
class FakeNet(nn.Module):
    def __init__(self, use_gpu=True, colorspace="Lab"):
        super(FakeNet, self).__init__()
        self.use_gpu = use_gpu
        self.colorspace = colorspace


class L2(FakeNet):
    def forward(self, in0, in1, retPerLayer=None):
        assert in0.size()[0] == 1  # currently only supports batchSize 1

        if self.colorspace == "RGB":
            (N, C, X, Y) = in0.size()
            value = torch.mean(
                torch.mean(torch.mean((in0 - in1) ** 2, dim=1).view(N, 1, X, Y), dim=2).view(N, 1, 1, Y), dim=3
            ).view(N)
            return value
        elif self.colorspace == "Lab":
            value = util.l2(
                util.tensor2np(util.tensor2tensorlab(in0.data, to_norm=False)),
                util.tensor2np(util.tensor2tensorlab(in1.data, to_norm=False)),
                range=100.0,
            ).astype("float")
            ret_var = Variable(torch.Tensor((value,)))
            if self.use_gpu:
                ret_var = ret_var.cuda()
            return ret_var


class DSSIM(FakeNet):
    def forward(self, in0, in1, retPerLayer=None):
        assert in0.size()[0] == 1  # currently only supports batchSize 1

        if self.colorspace == "RGB":
            value = util.dssim(1.0 * util.tensor2im(in0.data), 1.0 * util.tensor2im(in1.data), range=255.0).astype(
                "float"
            )
        elif self.colorspace == "Lab":
            value = util.dssim(
                util.tensor2np(util.tensor2tensorlab(in0.data, to_norm=False)),
                util.tensor2np(util.tensor2tensorlab(in1.data, to_norm=False)),
                range=100.0,
            ).astype("float")
        ret_var = Variable(torch.Tensor((value,)))
        if self.use_gpu:
            ret_var = ret_var.cuda()
        return ret_var


def print_network(net):
    num_params = 0
    for param in net.parameters():
        num_params += param.numel()
    print("Network", net)
    print("Total number of parameters: %d" % num_params)



================================================
FILE: validation/lpips/pretrained_networks.py
================================================
from collections import namedtuple
import torch
from torchvision import models as tv


class squeezenet(torch.nn.Module):
    def __init__(self, requires_grad=False, pretrained=True):
        super(squeezenet, self).__init__()
        pretrained_features = tv.squeezenet1_1(pretrained=pretrained).features
        self.slice1 = torch.nn.Sequential()
        self.slice2 = torch.nn.Sequential()
        self.slice3 = torch.nn.Sequential()
        self.slice4 = torch.nn.Sequential()
        self.slice5 = torch.nn.Sequential()
        self.slice6 = torch.nn.Sequential()
        self.slice7 = torch.nn.Sequential()
        self.N_slices = 7
        for x in range(2):
            self.slice1.add_module(str(x), pretrained_features[x])
        for x in range(2, 5):
            self.slice2.add_module(str(x), pretrained_features[x])
        for x in range(5, 8):
            self.slice3.add_module(str(x), pretrained_features[x])
        for x in range(8, 10):
            self.slice4.add_module(str(x), pretrained_features[x])
        for x in range(10, 11):
            self.slice5.add_module(str(x), pretrained_features[x])
        for x in range(11, 12):
            self.slice6.add_module(str(x), pretrained_features[x])
        for x in range(12, 13):
            self.slice7.add_module(str(x), pretrained_features[x])
        if not requires_grad:
            for param in self.parameters():
                param.requires_grad = False

    def forward(self, X):
        h = self.slice1(X)
        h_relu1 = h
        h = self.slice2(h)
        h_relu2 = h
        h = self.slice3(h)
        h_relu3 = h
        h = self.slice4(h)
        h_relu4 = h
        h = self.slice5(h)
        h_relu5 = h
        h = self.slice6(h)
        h_relu6 = h
        h = self.slice7(h)
        h_relu7 = h
        vgg_outputs = namedtuple("SqueezeOutputs", ["relu1", "relu2", "relu3", "relu4", "relu5", "relu6", "relu7"])
        out = vgg_outputs(h_relu1, h_relu2, h_relu3, h_relu4, h_relu5, h_relu6, h_relu7)

        return out


class alexnet(torch.nn.Module):
    def __init__(self, requires_grad=False, pretrained=True):
        super(alexnet, self).__init__()
        alexnet_pretrained_features = tv.alexnet(pretrained=pretrained).features
        self.slice1 = torch.nn.Sequential()
        self.slice2 = torch.nn.Sequential()
        self.slice3 = torch.nn.Sequential()
        self.slice4 = torch.nn.Sequential()
        self.slice5 = torch.nn.Sequential()
        self.N_slices = 5
        for x in range(2):
            self.slice1.add_module(str(x), alexnet_pretrained_features[x])
        for x in range(2, 5):
            self.slice2.add_module(str(x), alexnet_pretrained_features[x])
        for x in range(5, 8):
            self.slice3.add_module(str(x), alexnet_pretrained_features[x])
        for x in range(8, 10):
            self.slice4.add_module(str(x), alexnet_pretrained_features[x])
        for x in range(10, 12):
            self.slice5.add_module(str(x), alexnet_pretrained_features[x])
        if not requires_grad:
            for param in self.parameters():
                param.requires_grad = False

    def forward(self, X):
        h = self.slice1(X)
        h_relu1 = h
        h = self.slice2(h)
        h_relu2 = h
        h = self.slice3(h)
        h_relu3 = h
        h = self.slice4(h)
        h_relu4 = h
        h = self.slice5(h)
        h_relu5 = h
        alexnet_outputs = namedtuple("AlexnetOutputs", ["relu1", "relu2", "relu3", "relu4", "relu5"])
        out = alexnet_outputs(h_relu1, h_relu2, h_relu3, h_relu4, h_relu5)

        return out


class vgg16(torch.nn.Module):
    def __init__(self, requires_grad=False, pretrained=True):
        super(vgg16, self).__init__()
        vgg_pretrained_features = tv.vgg16(pretrained=pretrained).features
        self.slice1 = torch.nn.Sequential()
        self.slice2 = torch.nn.Sequential()
        self.slice3 = torch.nn.Sequential()
        self.slice4 = torch.nn.Sequential()
        self.slice5 = torch.nn.Sequential()
        self.N_slices = 5
        for x in range(4):
            self.slice1.add_module(str(x), vgg_pretrained_features[x])
        for x in range(4, 9):
            self.slice2.add_module(str(x), vgg_pretrained_features[x])
        for x in range(9, 16):
            self.slice3.add_module(str(x), vgg_pretrained_features[x])
        for x in range(16, 23):
            self.slice4.add_module(str(x), vgg_pretrained_features[x])
        for x in range(23, 30):
            self.slice5.add_module(str(x), vgg_pretrained_features[x])
        if not requires_grad:
            for param in self.parameters():
                param.requires_grad = False

    def forward(self, X):
        h = self.slice1(X)
        h_relu1_2 = h
        h = self.slice2(h)
        h_relu2_2 = h
        h = self.slice3(h)
        h_relu3_3 = h
        h = self.slice4(h)
        h_relu4_3 = h
        h = self.slice5(h)
        h_relu5_3 = h
        vgg_outputs = namedtuple("VggOutputs", ["relu1_2", "relu2_2", "relu3_3", "relu4_3", "relu5_3"])
        out = vgg_outputs(h_relu1_2, h_relu2_2, h_relu3_3, h_relu4_3, h_relu5_3)

        return out


class resnet(torch.nn.Module):
    def __init__(self, requires_grad=False, pretrained=True, num=18):
        super(resnet, self).__init__()
        if num == 18:
            self.net = tv.resnet18(pretrained=pretrained)
        elif num == 34:
            self.net = tv.resnet34(pretrained=pretrained)
        elif num == 50:
            self.net = tv.resnet50(pretrained=pretrained)
        elif num == 101:
            self.net = tv.resnet101(pretrained=pretrained)
        elif num == 152:
            self.net = tv.resnet152(pretrained=pretrained)
        self.N_slices = 5

        self.conv1 = self.net.conv1
        self.bn1 = self.net.bn1
        self.relu = self.net.relu
        self.maxpool = self.net.maxpool
        self.layer1 = self.net.layer1
        self.layer2 = self.net.layer2
        self.layer3 = self.net.layer3
        self.layer4 = self.net.layer4

    def forward(self, X):
        h = self.conv1(X)
        h = self.bn1(h)
        h = self.relu(h)
        h_relu1 = h
        h = self.maxpool(h)
        h = self.layer1(h)
        h_conv2 = h
        h = self.layer2(h)
        h_conv3 = h
        h = self.layer3(h)
        h_conv4 = h
        h = self.layer4(h)
        h_conv5 = h

        outputs = namedtuple("Outputs", ["relu1", "conv2", "conv3", "conv4", "conv5"])
        out = outputs(h_relu1, h_conv2, h_conv3, h_conv4, h_conv5)

        return out


================================================
FILE: validation/lpips/util.py
================================================
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
from skimage.measure import compare_ssim
import torch


def normalize_tensor(in_feat, eps=1e-10):
    norm_factor = torch.sqrt(torch.sum(in_feat ** 2, dim=1, keepdim=True))
    return in_feat / (norm_factor + eps)


def l2(p0, p1, range=255.0):
    return 0.5 * np.mean((p0 / range - p1 / range) ** 2)


def psnr(p0, p1, peak=255.0):
    return 10 * np.log10(peak ** 2 / np.mean((1.0 * p0 - 1.0 * p1) ** 2))


def dssim(p0, p1, range=255.0):
    return (1 - compare_ssim(p0, p1, data_range=range, multichannel=True)) / 2.0


def rgb2lab(in_img, mean_cent=False):
    from skimage import color

    img_lab = color.rgb2lab(in_img)
    if mean_cent:
        img_lab[:, :, 0] = img_lab[:, :, 0] - 50
    return img_lab


def tensor2np(tensor_obj):
    # change dimension of a tensor object into a numpy array
    return tensor_obj[0].cpu().float().numpy().transpose((1, 2, 0))


def np2tensor(np_obj):
    # change dimenion of np array into tensor array
    return torch.Tensor(np_obj[:, :, :, np.newaxis].transpose((3, 2, 0, 1)))


def tensor2tensorlab(image_tensor, to_norm=True, mc_only=False):
    # image tensor to lab tensor
    from skimage import color

    img = tensor2im(image_tensor)
    img_lab = color.rgb2lab(img)
    if mc_only:
        img_lab[:, :, 0] = img_lab[:, :, 0] - 50
    if to_norm and not mc_only:
        img_lab[:, :, 0] = img_lab[:, :, 0] - 50
        img_lab = img_lab / 100.0

    return np2tensor(img_lab)


def tensorlab2tensor(lab_tensor, return_inbnd=False):
    from skimage import color
    import warnings

    warnings.filterwarnings("ignore")

    lab = tensor2np(lab_tensor) * 100.0
    lab[:, :, 0] = lab[:, :, 0] + 50

    rgb_back = 255.0 * np.clip(color.lab2rgb(lab.astype("float")), 0, 1)
    if return_inbnd:
        # convert back to lab, see if we match
        lab_back = color.rgb2lab(rgb_back.astype("uint8"))
        mask = 1.0 * np.isclose(lab_back, lab, atol=2.0)
        mask = np2tensor(np.prod(mask, axis=2)[:, :, np.newaxis])
        return (im2tensor(rgb_back), mask)
    else:
        return im2tensor(rgb_back)


def rgb2lab(input):
    from skimage import color

    return color.rgb2lab(input / 255.0)


def tensor2im(image_tensor, imtype=np.uint8, cent=1.0, factor=255.0 / 2.0):
    image_numpy = image_tensor[0].cpu().float().numpy()
    image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + cent) * factor
    return image_numpy.astype(imtype)


def im2tensor(image, imtype=np.uint8, cent=1.0, factor=255.0 / 2.0):
    return torch.Tensor((image / factor - cent)[:, :, :, np.newaxis].transpose((3, 2, 0, 1)))


def tensor2vec(vector_tensor):
    return vector_tensor.data.cpu().numpy()[:, :, 0, 0]


def voc_ap(rec, prec, use_07_metric=False):
    """ ap = voc_ap(rec, prec, [use_07_metric])
    Compute VOC AP given precision and recall.
    If use_07_metric is true, uses the
    VOC 07 11 point method (default:False).
    """
    if use_07_metric:
        # 11 point metric
        ap = 0.0
        for t in np.arange(0.0, 1.1, 0.1):
            if np.sum(rec >= t) == 0:
                p = 0
            else:
                p = np.max(prec[rec >= t])
            ap = ap + p / 11.0
    else:
        # correct AP calculation
        # first append sentinel values at the end
        mrec = np.concatenate(([0.0], rec, [1.0]))
        mpre = np.concatenate(([0.0], prec, [0.0]))

        # compute the precision envelope
        for i in range(mpre.size - 1, 0, -1):
            mpre[i - 1] = np.maximum(mpre[i - 1], mpre[i])

        # to calculate area under PR curve, look for points
        # where X axis (recall) changes value
        i = np.where(mrec[1:] != mrec[:-1])[0]

        # and sum (\Delta recall) * prec
        ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1])
    return ap


def tensor2im(image_tensor, imtype=np.uint8, cent=1.0, factor=255.0 / 2.0):
    # def tensor2im(image_tensor, imtype=np.uint8, cent=1., factor=1.):
    image_numpy = image_tensor[0].cpu().float().numpy()
    image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + cent) * factor
    return image_numpy.astype(imtype)


def im2tensor(image, imtype=np.uint8, cent=1.0, factor=255.0 / 2.0):
    # def im2tensor(image, imtype=np.uint8, cent=1., factor=1.):
    return torch.Tensor((image / factor - cent)[:, :, :, np.newaxis].transpose((3, 2, 0, 1)))



================================================
FILE: validation/metrics.py
================================================
import os
import pickle
import random

from sklearn.metrics import pairwise_distances
from tqdm import tqdm
import torch
from torch.nn import functional as F
import numpy as np
from scipy import linalg

from .inception import InceptionV3
from . import lpips


@torch.no_grad()
def vae_fid(vae, batch_size, latent_dim, n_sample, inception_name, calculate_prdc=True):
    vae.eval()

    inception = InceptionV3([3], normalize_input=False, init_weights=False)
    inception = inception.eval().to(next(vae.parameters()).device)

    n_batch = n_sample // batch_size
    resid = n_sample - (n_batch * batch_size)
    if resid == 0:
        batch_sizes = [batch_size] * n_batch
    else:
        batch_sizes = [batch_size] * n_batch + [resid]
    features = []

    for batch in batch_sizes:
        latent = torch.randn(batch, *latent_dim).cuda()
        img = vae.decode(latent)
        feat = inception(img)[0].view(img.shape[0], -1)
        features.append(feat.to("cpu"))
    features = torch.cat(features, 0).numpy()

    del inception

    sample_mean = np.mean(features, 0)
    sample_cov = np.cov(features, rowvar=False)

    with open(f"inception_{inception_name}_stats.pkl", "rb") as f:
        embeds = pickle.load(f)
        real_mean = embeds["mean"]
        real_cov = embeds["cov"]

    cov_sqrt, _ = linalg.sqrtm(sample_cov @ real_cov, disp=False)

    if not np.isfinite(cov_sqrt).all():
        print("product of cov matrices is singular")
        offset = np.eye(sample_cov.shape[0]) * 1e-6
        cov_sqrt = linalg.sqrtm((sample_cov + offset) @ (real_cov + offset))

    if np.iscomplexobj(cov_sqrt):
        if not np.allclose(np.diagonal(cov_sqrt).imag, 0, atol=1e-3):
            m = np.max(np.abs(cov_sqrt.imag))

            raise ValueError(f"Imaginary component {m}")

        cov_sqrt = cov_sqrt.real

    mean_diff = sample_mean - real_mean
    mean_norm = mean_diff @ mean_diff

    trace = np.trace(sample_cov) + np.trace(real_cov) - 2 * np.trace(cov_sqrt)

    fid = mean_norm + trace

    ret_dict = {"FID": fid}

    if calculate_prdc:
        with open(f"inception_{inception_name}_features.pkl", "rb") as f:
            embeds = pickle.load(f)
            real_feats = embeds["features"]
        _, _, density, coverage = prdc(real_feats[:80000], features[:80000])
        ret_dict["Density"] = density
        ret_dict["Coverage"] = coverage

    return ret_dict


@torch.no_grad()
def fid(generator, batch_size, n_sample, truncation, inception_name, calculate_prdc=True):
    generator.eval()
    mean_latent = generator.mean_latent(2 ** 14)

    inception = InceptionV3([3], normalize_input=False, init_weights=False)
    inception = inception.eval().to(next(generator.parameters()).device)

    n_batch = n_sample // batch_size
    resid = n_sample - (n_batch * batch_size)
    if resid == 0:
        batch_sizes = [batch_size] * n_batch
    else:
        batch_sizes = [batch_size] * n_batch + [resid]
    features = []

    for batch in batch_sizes:
        if truncation is None:
            trunc = random.uniform(0.9, 1.5)
        else:
            trunc = truncation
        latent = torch.randn(batch, 512).cuda()
        img, _ = generator([latent], truncation=trunc, truncation_latent=mean_latent)
        feat = inception(img)[0].view(img.shape[0], -1)
        features.append(feat.to("cpu"))
    features = torch.cat(features, 0).numpy()

    del inception

    sample_mean = np.mean(features, 0)
    sample_cov = np.cov(features, rowvar=False)

    with open(f"inception_{inception_name}_stats.pkl", "rb") as f:
        embeds = pickle.load(f)
        real_mean = embeds["mean"]
        real_cov = embeds["cov"]

    cov_sqrt, _ = linalg.sqrtm(sample_cov @ real_cov, disp=False)

    if not np.isfinite(cov_sqrt).all():
        print("product of cov matrices is singular")
        offset = np.eye(sample_cov.shape[0]) * 1e-6
        cov_sqrt = linalg.sqrtm((sample_cov + offset) @ (real_cov + offset))

    if np.iscomplexobj(cov_sqrt):
        if not np.allclose(np.diagonal(cov_sqrt).imag, 0, atol=1e-3):
            m = np.max(np.abs(cov_sqrt.imag))

            raise ValueError(f"Imaginary component {m}")

        cov_sqrt = cov_sqrt.real

    mean_diff = sample_mean - real_mean
    mean_norm = mean_diff @ mean_diff

    trace = np.trace(sample_cov) + np.trace(real_cov) - 2 * np.trace(cov_sqrt)

    fid = mean_norm + trace

    ret_dict = {"FID": fid}

    if calculate_prdc:
        with open(f"inception_{inception_name}_features.pkl", "rb") as f:
            embeds = pickle.load(f)
            real_feats = embeds["features"]
        _, _, density, coverage = prdc(real_feats[:80000], features[:80000])
        ret_dict["Density"] = density
        ret_dict["Coverage"] = coverage

    return ret_dict


def get_dataset_inception_features(loader, inception_name, size):
    if not os.path.exists(f"inception_{inception_name}_stats.pkl"):
        print("calculating inception features for FID....")
        inception = InceptionV3([3], normalize_input=False, init_weights=False)
        inception = torch.nn.DataParallel(inception).eval().cuda()

        feature_list = []
        for img in tqdm(loader):
            img = img.cuda()
            feature = inception(img)[0].view(img.shape[0], -1)
            feature_list.append(feature.to("cpu"))
        features = torch.cat(feature_list, 0).numpy()

        mean = np.mean(features, 0)
        cov = np.cov(features, rowvar=False)

        with open(f"inception_{inception_name}_stats.pkl", "wb") as f:
            pickle.dump({"mean": mean, "cov": cov, "size": size, "feat": features}, f)
        with open(f"inception_{inception_name}_features.pkl", "wb") as f:
            pickle.dump({"features": features}, f)
    else:
        print(f"Found inception features: inception_{inception_name}_stats.pkl")


def compute_pairwise_distance(data_x, data_y=None, metric="l2"):
    if data_y is None:
        data_y = data_x
    dists = pairwise_distances(
        data_x.reshape((len(data_x), -1)), data_y.reshape((len(data_y), -1)), metric=metric, n_jobs=24
    )
    return dists


def get_kth_value(unsorted, k, axis=-1):
    indices = np.argpartition(unsorted, k, axis=axis)[..., :k]
    k_smallests = np.take_along_axis(unsorted, indices, axis=axis)
    kth_values = k_smallests.max(axis=axis)
    return kth_values


def compute_nearest_neighbour_distances(input_features, nearest_k, metric):
    distances = compute_pairwise_distance(input_features, metric=metric)
    radii = get_kth_value(distances, k=nearest_k + 1, axis=-1)
    return radii


def prdc(real_features, fake_features, nearest_k=10, metric="l2"):
    real_nearest_neighbour_distances = compute_nearest_neighbour_distances(real_features, nearest_k, metric=metric)
    fake_nearest_neighbour_distances = compute_nearest_neighbour_distances(fake_features, nearest_k, metric=metric)
    distance_real_fake = compute_pairwise_distance(real_features, fake_features, metric=metric)

    precision = (distance_real_fake < np.expand_dims(real_nearest_neighbour_distances, axis=1)).any(axis=0).mean()
    recall = (distance_real_fake < np.expand_dims(fake_nearest_neighbour_distances, axis=0)).any(axis=1).mean()

    density = (1.0 / float(nearest_k)) * (
        distance_real_fake < np.expand_dims(real_nearest_neighbour_distances, axis=1)
    ).sum(axis=0).mean()
    coverage = (distance_real_fake.min(axis=1) < real_nearest_neighbour_distances).mean()

    return precision, recall, density, coverage


def lerp(a, b, t):
    return a + (b - a) * t


@torch.no_grad()
def ppl(generator, batch_size, n_sample, space, crop, latent_dim, eps=1e-4):
    generator.eval()

    percept = lpips.PerceptualLoss(
        model="net-lin", net="vgg", use_gpu=True, gpu_ids=[next(generator.parameters()).device.index]
    )

    distances = []

    n_batch = n_sample // batch_size
    resid = n_sample - (n_batch * batch_size)
    if resid == 0:
        batch_sizes = [batch_size] * n_batch
    else:
        batch_sizes = [batch_size] * n_batch + [resid]

    for batch_size in batch_sizes:
        noise = generator.make_noise()

        inputs = torch.randn([batch_size * 2, latent_dim]).cuda()
        lerp_t = torch.rand(batch_size).cuda()

        if space == "w":
            latent = generator.get_latent(inputs)
            latent_t0, latent_t1 = latent[::2], latent[1::2]
            latent_e0 = lerp(latent_t0, latent_t1, lerp_t[:, None])
            latent_e1 = lerp(latent_t0, latent_t1, lerp_t[:, None] + eps)
            latent_e = torch.stack([latent_e0, latent_e1], 1).view(*latent.shape)

        image, _ = generator(latent_e, input_is_latent=True, noise=noise)

        if crop:
            c = image.shape[2] // 8
            image = image[:, :, c * 3 : c * 7, c * 2 : c * 6]

        factor = image.shape[2] // 256

        if factor > 1:
            image = F.interpolate(image, size=(256, 256), mode="bilinear", align_corners=False)

        dist = percept(image[::2], image[1::2]).view(image.shape[0] // 2) / (eps ** 2)
        distances.append(dist.to("cpu").numpy())

    distances = np.concatenate(distances, 0)

    lo = np.percentile(distances, 1, interpolation="lower")
    hi = np.percentile(distances, 99, interpolation="higher")
    filtered_dist = np.extract(np.logical_and(lo <= distances, distances <= hi), distances)
    path_length = filtered_dist.mean()

    del percept, inputs, lerp_t, image, dist

    return path_length


================================================
FILE: validation/spectral_norm.py
================================================
import torch


class SpectralNorm(object):
    def __init__(self, name="weight", n_power_iterations=1, dim=0, eps=1e-12):
        self.name = name
        self.dim = dim
        if n_power_iterations <= 0:
            raise ValueError(
                "Expected n_power_iterations to be positive, but "
                "got n_power_iterations={}".format(n_power_iterations)
            )
        self.n_power_iterations = n_power_iterations
        self.eps = eps

    def reshape_weight_to_matrix(self, weight):
        weight_mat = weight
        if self.dim != 0:
            # permute dim to front
            weight_mat = weight_mat.permute(self.dim, *[d for d in range(weight_mat.dim()) if d != self.dim])
        height = weight_mat.size(0)
        return weight_mat.reshape(height, -1)

    def compute_sigma(self, module):
        with torch.no_grad():
            weight = getattr(module, self.name)
            weight_mat = self.reshape_weight_to_matrix(weight)

            u = getattr(module, self.name + "_u")
            v = getattr(module, self.name + "_v")
            for _ in range(self.n_power_iterations):
                v = torch.nn.functional.normalize(torch.mv(weight_mat.t(), u), dim=0, eps=self.eps)
                u = torch.nn.functional.normalize(torch.mv(weight_mat, v), dim=0, eps=self.eps)
            setattr(module, self.name + "_u", u)
            setattr(module, self.name + "_v", v)

            sigma = torch.dot(u, torch.mv(weight_mat, v))
            setattr(module, "spectral_norm", sigma)

    def remove(self, module):
        delattr(module, self.name)
        delattr(module, self.name + "_u")
        delattr(module, self.name + "_v")
        delattr(module, "spectral_norm")

    def __call__(self, module, inputs):
        self.compute_sigma(module)

    def _solve_v_and_rescale(self, weight_mat, u, target_sigma):
        v = torch.chain_matmul(weight_mat.t().mm(weight_mat).pinverse(), weight_mat.t(), u.unsqueeze(1)).squeeze(1)
        return v.mul_(target_sigma / torch.dot(u, torch.mv(weight_mat, v)))

    @staticmethod
    def apply(module, name, n_power_iterations, dim, eps, normalize=True):
        for k, hook in module._forward_pre_hooks.items():
            if isinstance(hook, SpectralNorm) and hook.name == name:
                raise RuntimeError("Cannot register two spectral_norm hooks on " "the same parameter {}".format(name))

        fn = SpectralNorm(name, n_power_iterations, dim, eps)
        weight = module._parameters[name]
        with torch.no_grad():
            weight_mat = fn.reshape_weight_to_matrix(weight)

            h, w = weight_mat.size()
            u = torch.nn.functional.normalize(weight.new_empty(h).normal_(0, 1), dim=0, eps=fn.eps)
            v = torch.nn.functional.normalize(weight.new_empty(w).normal_(0, 1), dim=0, eps=fn.eps)

        module.register_buffer(fn.name + "_u", u)
        module.register_buffer(fn.name + "_v", v)
        module.register_buffer("spectral_norm", torch.tensor(-1, device=next(module.parameters()).device))

        module.register_forward_pre_hook(fn)
        return fn


def track_spectral_norm(module, name="weight", n_power_iterations=1, eps=1e-12, dim=None):
    r"""Tracks the spectral norm of a module's weight parameter
    Args:
        module (nn.Module): containing module
        name (str, optional): name of weight parameter
        n_power_iterations (int, optional): number of power iterations to
            calculate spectral norm
        eps (float, optional): epsilon for numerical stability in
            calculating norms
        dim (int, optional): dimension corresponding to number of outputs,
            the default is ``0``, except for modules that are instances of
            ConvTranspose{1,2,3}d, when it is ``1``
    Returns:
        The original module with the spectral norm hook
    Example::
        >>> m = spectral_norm(nn.Linear(20, 40))
        >>> m
        Linear(in_features=20, out_features=40, bias=True)
        >>> m.weight_u.size()
        torch.Size([40])
    """
    if dim is None:
        if isinstance(module, (torch.nn.ConvTranspose1d, torch.nn.ConvTranspose2d, torch.nn.ConvTranspose3d)):
            dim = 1
        else:
            dim = 0
    SpectralNorm.apply(module, name, n_power_iterations, dim, eps)
    return module


def remove_spectral_norm(module, name="weight"):
    r"""Removes the spectral normalization reparameterization from a module.
    Args:
        module (Module): containing module
        name (str, optional): name of weight parameter
    Example:
        >>> m = spectral_norm(nn.Linear(40, 10))
        >>> remove_spectral_norm(m)
    """
    for k, hook in module._forward_pre_hooks.items():
        if isinstance(hook, SpectralNorm) and hook.name == name:
            hook.remove(module)
            del module._forward_pre_hooks[k]
            break
    else:
        raise ValueError("spectral_norm of '{}' not found in {}".format(name, module))

    return module


================================================
FILE: workspace/naamloos_metadata.json
================================================
{"total_frames": 4986}


================================================
FILE: workspace/naamloos_params.json
================================================
{"intro_num_beats": 64, "intro_loop_smoothing": 30, "intro_loop_factor": 0.4, "intro_loop_len": 12, "drop_num_beats": 32, "drop_loop_smoothing": 15, "drop_loop_factor": 1, "drop_loop_len": 6, "onset_smooth": 2, "onset_clip": 95, "freq_mod": 10, "freq_mod_offset": 0, "freq_smooth": 5, "freq_latent_smooth": 4, "freq_latent_layer": 1, "freq_latent_weight": 2, "high_freq_mod": 10, "high_freq_mod_offset": 0, "high_freq_smooth": 4, "high_freq_latent_smooth": 5, "high_freq_latent_layer": 2, "high_freq_latent_weight": 1.5, "rms_smooth": 5, "bass_smooth": 5, "bass_clip": 65, "drop_clip": 75, "drop_smooth": 5, "drop_weight": 1, "high_noise_clip": 100, "high_noise_weight": 1.5, "low_noise_weight": 1}