[
  {
    "path": ".gitignore",
    "content": "pretrained_models/\nwandb\nwandb/\n*.lmdb/\n*.pkl\ncheckpoints/\nmaua-stylegan/\n.vscode\noutput/\nworkspace/*\n!workspace\noutput/*\n!output\n\n# Byte-compiled / optimized / DLL files\n__pycache__/\n*.py[cod]\n*$py.class\n\n# C extensions\n*.so\n\n# Distribution / packaging\n.Python\nbuild/\ndevelop-eggs/\ndist/\ndownloads/\neggs/\n.eggs/\nlib/\nlib64/\nparts/\nsdist/\nvar/\nwheels/\npip-wheel-metadata/\nshare/python-wheels/\n*.egg-info/\n.installed.cfg\n*.egg\nMANIFEST\n\n# PyInstaller\n#  Usually these files are written by a python script from a template\n#  before PyInstaller builds the exe, so as to inject date/other infos into it.\n*.manifest\n*.spec\n\n# Installer logs\npip-log.txt\npip-delete-this-directory.txt\n\n# Unit test / coverage reports\nhtmlcov/\n.tox/\n.nox/\n.coverage\n.coverage.*\n.cache\nnosetests.xml\ncoverage.xml\n*.cover\n*.py,cover\n.hypothesis/\n.pytest_cache/\n\n# Translations\n*.mo\n*.pot\n\n# Django stuff:\n*.log\nlocal_settings.py\ndb.sqlite3\ndb.sqlite3-journal\n\n# Flask stuff:\ninstance/\n.webassets-cache\n\n# Scrapy stuff:\n.scrapy\n\n# Sphinx documentation\ndocs/_build/\n\n# PyBuilder\ntarget/\n\n# Jupyter Notebook\n.ipynb_checkpoints\n\n# IPython\nprofile_default/\nipython_config.py\n\n# pyenv\n.python-version\n\n# pipenv\n#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.\n#   However, in case of collaboration, if having platform-specific dependencies or dependencies\n#   having no cross-platform support, pipenv may install dependencies that don't work, or not\n#   install all needed dependencies.\n#Pipfile.lock\n\n# PEP 582; used by e.g. github.com/David-OConnor/pyflow\n__pypackages__/\n\n# Celery stuff\ncelerybeat-schedule\ncelerybeat.pid\n\n# SageMath parsed files\n*.sage.py\n\n# Environments\n.env\n.venv\nenv/\nvenv/\nENV/\nenv.bak/\nvenv.bak/\n\n# Spyder project settings\n.spyderproject\n.spyproject\n\n# Rope project settings\n.ropeproject\n\n# mkdocs documentation\n/site\n\n# mypy\n.mypy_cache/\n.dmypy.json\ndmypy.json\n\n# Pyre type checker\n.pyre/\n\n"
  },
  {
    "path": "LICENSE/LICENSE-AUDIOREACTIVE",
    "content": "Code for Audio-reactive Latent Interpolations with StyleGAN\nIncluding the folder audioreactive/, generate_audiovisual.py, generate_video.py, select_latents.py, and render.py\n\nCopyright (C) 2020 Hans Brouwer\n\nThis program is free software: you can redistribute it and/or modify\nit under the terms of the GNU General Public License as published by\nthe Free Software Foundation, either version 3 of the License, or\n(at your option) any later version.\n\nThis program is distributed in the hope that it will be useful,\nbut WITHOUT ANY WARRANTY; without even the implied warranty of\nMERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the\nGNU General Public License for more details.\n\nTERMS AND CONDITIONS\n0. Definitions.\n\n“This License” refers to version 3 of the GNU General Public License.\n\n“Copyright” also means copyright-like laws that apply to other kinds of works, such as semiconductor masks.\n\n“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.\n\nTo “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.\n\nA “covered work” means either the unmodified Program or a work based on the Program.\n\nTo “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.\n\nTo “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.\n\nAn 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.\n1. Source Code.\n\nThe “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.\n\nA “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.\n\nThe “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.\n\nThe “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.\n\nThe Corresponding Source need not include anything that users can regenerate automatically from other parts of the Corresponding Source.\n\nThe Corresponding Source for a work in source code form is that same work.\n2. Basic Permissions.\n\nAll 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.\n\nYou 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.\n\nConveying under any other circumstances is permitted solely under the conditions stated below. Sublicensing is not allowed; section 10 makes it unnecessary.\n3. Protecting Users' Legal Rights From Anti-Circumvention Law.\n\nNo 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.\n\nWhen 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.\n4. Conveying Verbatim Copies.\n\nYou 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.\n\nYou may charge any price or no price for each copy that you convey, and you may offer support or warranty protection for a fee.\n5. Conveying Modified Source Versions.\n\nYou 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:\n\n    a) The work must carry prominent notices stating that you modified it, and giving a relevant date.\n    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”.\n    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.\n    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.\n\nA 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.\n6. Conveying Non-Source Forms.\n\nYou 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:\n\n    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.\n    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.\n    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.\n    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.\n    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.\n\nA 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.\n\nA “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.\n\n“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.\n\nIf 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).\n\nThe 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.\n\nCorresponding 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.\n7. Additional Terms.\n\n“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.\n\nWhen 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.\n\nNotwithstanding 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:\n\n    a) Disclaiming warranty or limiting liability differently from the terms of sections 15 and 16 of this License; or\n    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\n    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\n    d) Limiting the use for publicity purposes of names of licensors or authors of the material; or\n    e) Declining to grant rights under trademark law for use of some trade names, trademarks, or service marks; or\n    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.\n\nAll 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.\n\nIf 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.\n\nAdditional 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.\n8. Termination.\n\nYou 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).\n\nHowever, 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.\n\nMoreover, 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.\n\nTermination 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.\n9. Acceptance Not Required for Having Copies.\n\nYou 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.\n10. Automatic Licensing of Downstream Recipients.\n\nEach 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.\n\nAn “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.\n\nYou 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.\n11. Patents.\n\nA “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”.\n\nA 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.\n\nEach 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.\n\nIn 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.\n\nIf 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.\n\nIf, 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.\n\nA 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.\n\nNothing 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.\n12. No Surrender of Others' Freedom.\n\nIf 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.\n13. Use with the GNU Affero General Public License.\n\nNotwithstanding 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.\n14. Revised Versions of this License.\n\nThe 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.\n\nEach 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.\n\nIf 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.\n\nLater 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.\n15. Disclaimer of Warranty.\n\nTHERE 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.\n16. Limitation of Liability.\n\nIN 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.\n17. Interpretation of Sections 15 and 16.\n\nIf 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.\n\nEND OF TERMS AND CONDITIONS"
  },
  {
    "path": "LICENSE/LICENSE-AUTOENCODER",
    "content": "                                 Apache License\n                           Version 2.0, January 2004\n                        http://www.apache.org/licenses/\n\n   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION\n\n   1. Definitions.\n\n      \"License\" shall mean the terms and conditions for use, reproduction,\n      and distribution as defined by Sections 1 through 9 of this document.\n\n      \"Licensor\" shall mean the copyright owner or entity authorized by\n      the copyright owner that is granting the License.\n\n      \"Legal Entity\" shall mean the union of the acting entity and all\n      other entities that control, are controlled by, or are under common\n      control with that entity. For the purposes of this definition,\n      \"control\" means (i) the power, direct or indirect, to cause the\n      direction or management of such entity, whether by contract or\n      otherwise, or (ii) ownership of fifty percent (50%) or more of the\n      outstanding shares, or (iii) beneficial ownership of such entity.\n\n      \"You\" (or \"Your\") shall mean an individual or Legal Entity\n      exercising permissions granted by this License.\n\n      \"Source\" form shall mean the preferred form for making modifications,\n      including but not limited to software source code, documentation\n      source, and configuration files.\n\n      \"Object\" form shall mean any form resulting from mechanical\n      transformation or translation of a Source form, including but\n      not limited to compiled object code, generated documentation,\n      and conversions to other media types.\n\n      \"Work\" shall mean the work of authorship, whether in Source or\n      Object form, made available under the License, as indicated by a\n      copyright notice that is included in or attached to the work\n      (an example is provided in the Appendix below).\n\n      \"Derivative Works\" shall mean any work, whether in Source or Object\n      form, that is based on (or derived from) the Work and for which the\n      editorial revisions, annotations, elaborations, or other modifications\n      represent, as a whole, an original work of authorship. For the purposes\n      of this License, Derivative Works shall not include works that remain\n      separable from, or merely link (or bind by name) to the interfaces of,\n      the Work and Derivative Works thereof.\n\n      \"Contribution\" shall mean any work of authorship, including\n      the original version of the Work and any modifications or additions\n      to that Work or Derivative Works thereof, that is intentionally\n      submitted to Licensor for inclusion in the Work by the copyright owner\n      or by an individual or Legal Entity authorized to submit on behalf of\n      the copyright owner. For the purposes of this definition, \"submitted\"\n      means any form of electronic, verbal, or written communication sent\n      to the Licensor or its representatives, including but not limited to\n      communication on electronic mailing lists, source code control systems,\n      and issue tracking systems that are managed by, or on behalf of, the\n      Licensor for the purpose of discussing and improving the Work, but\n      excluding communication that is conspicuously marked or otherwise\n      designated in writing by the copyright owner as \"Not a Contribution.\"\n\n      \"Contributor\" shall mean Licensor and any individual or Legal Entity\n      on behalf of whom a Contribution has been received by Licensor and\n      subsequently incorporated within the Work.\n\n   2. Grant of Copyright License. Subject to the terms and conditions of\n      this License, each Contributor hereby grants to You a perpetual,\n      worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n      copyright license to reproduce, prepare Derivative Works of,\n      publicly display, publicly perform, sublicense, and distribute the\n      Work and such Derivative Works in Source or Object form.\n\n   3. Grant of Patent License. Subject to the terms and conditions of\n      this License, each Contributor hereby grants to You a perpetual,\n      worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n      (except as stated in this section) patent license to make, have made,\n      use, offer to sell, sell, import, and otherwise transfer the Work,\n      where such license applies only to those patent claims licensable\n      by such Contributor that are necessarily infringed by their\n      Contribution(s) alone or by combination of their Contribution(s)\n      with the Work to which such Contribution(s) was submitted. If You\n      institute patent litigation against any entity (including a\n      cross-claim or counterclaim in a lawsuit) alleging that the Work\n      or a Contribution incorporated within the Work constitutes direct\n      or contributory patent infringement, then any patent licenses\n      granted to You under this License for that Work shall terminate\n      as of the date such litigation is filed.\n\n   4. Redistribution. You may reproduce and distribute copies of the\n      Work or Derivative Works thereof in any medium, with or without\n      modifications, and in Source or Object form, provided that You\n      meet the following conditions:\n\n      (a) You must give any other recipients of the Work or\n          Derivative Works a copy of this License; and\n\n      (b) You must cause any modified files to carry prominent notices\n          stating that You changed the files; and\n\n      (c) You must retain, in the Source form of any Derivative Works\n          that You distribute, all copyright, patent, trademark, and\n          attribution notices from the Source form of the Work,\n          excluding those notices that do not pertain to any part of\n          the Derivative Works; and\n\n      (d) If the Work includes a \"NOTICE\" text file as part of its\n          distribution, then any Derivative Works that You distribute must\n          include a readable copy of the attribution notices contained\n          within such NOTICE file, excluding those notices that do not\n          pertain to any part of the Derivative Works, in at least one\n          of the following places: within a NOTICE text file distributed\n          as part of the Derivative Works; within the Source form or\n          documentation, if provided along with the Derivative Works; or,\n          within a display generated by the Derivative Works, if and\n          wherever such third-party notices normally appear. The contents\n          of the NOTICE file are for informational purposes only and\n          do not modify the License. You may add Your own attribution\n          notices within Derivative Works that You distribute, alongside\n          or as an addendum to the NOTICE text from the Work, provided\n          that such additional attribution notices cannot be construed\n          as modifying the License.\n\n      You may add Your own copyright statement to Your modifications and\n      may provide additional or different license terms and conditions\n      for use, reproduction, or distribution of Your modifications, or\n      for any such Derivative Works as a whole, provided Your use,\n      reproduction, and distribution of the Work otherwise complies with\n      the conditions stated in this License.\n\n   5. Submission of Contributions. Unless You explicitly state otherwise,\n      any Contribution intentionally submitted for inclusion in the Work\n      by You to the Licensor shall be under the terms and conditions of\n      this License, without any additional terms or conditions.\n      Notwithstanding the above, nothing herein shall supersede or modify\n      the terms of any separate license agreement you may have executed\n      with Licensor regarding such Contributions.\n\n   6. Trademarks. This License does not grant permission to use the trade\n      names, trademarks, service marks, or product names of the Licensor,\n      except as required for reasonable and customary use in describing the\n      origin of the Work and reproducing the content of the NOTICE file.\n\n   7. Disclaimer of Warranty. Unless required by applicable law or\n      agreed to in writing, Licensor provides the Work (and each\n      Contributor provides its Contributions) on an \"AS IS\" BASIS,\n      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or\n      implied, including, without limitation, any warranties or conditions\n      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A\n      PARTICULAR PURPOSE. You are solely responsible for determining the\n      appropriateness of using or redistributing the Work and assume any\n      risks associated with Your exercise of permissions under this License.\n\n   8. Limitation of Liability. In no event and under no legal theory,\n      whether in tort (including negligence), contract, or otherwise,\n      unless required by applicable law (such as deliberate and grossly\n      negligent acts) or agreed to in writing, shall any Contributor be\n      liable to You for damages, including any direct, indirect, special,\n      incidental, or consequential damages of any character arising as a\n      result of this License or out of the use or inability to use the\n      Work (including but not limited to damages for loss of goodwill,\n      work stoppage, computer failure or malfunction, or any and all\n      other commercial damages or losses), even if such Contributor\n      has been advised of the possibility of such damages.\n\n   9. Accepting Warranty or Additional Liability. While redistributing\n      the Work or Derivative Works thereof, You may choose to offer,\n      and charge a fee for, acceptance of support, warranty, indemnity,\n      or other liability obligations and/or rights consistent with this\n      License. However, in accepting such obligations, You may act only\n      on Your own behalf and on Your sole responsibility, not on behalf\n      of any other Contributor, and only if You agree to indemnify,\n      defend, and hold each Contributor harmless for any liability\n      incurred by, or claims asserted against, such Contributor by reason\n      of your accepting any such warranty or additional liability.\n\n   END OF TERMS AND CONDITIONS\n\n   APPENDIX: How to apply the Apache License to your work.\n\n      To apply the Apache License to your work, attach the following\n      boilerplate notice, with the fields enclosed by brackets \"[]\"\n      replaced with your own identifying information. (Don't include\n      the brackets!)  The text should be enclosed in the appropriate\n      comment syntax for the file format. We also recommend that a\n      file or class name and description of purpose be included on the\n      same \"printed page\" as the copyright notice for easier\n      identification within third-party archives.\n\n   Copyright [yyyy] [name of copyright owner]\n\n   Licensed under the Apache License, Version 2.0 (the \"License\");\n   you may not use this file except in compliance with the License.\n   You may obtain a copy of the License at\n\n       http://www.apache.org/licenses/LICENSE-2.0\n\n   Unless required by applicable law or agreed to in writing, software\n   distributed under the License is distributed on an \"AS IS\" BASIS,\n   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n   See the License for the specific language governing permissions and\n   limitations under the License.\n"
  },
  {
    "path": "LICENSE/LICENSE-CONTRASTIVE-LEARNER",
    "content": "MIT License\n\nCopyright (c) 2020 Phil Wang\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n"
  },
  {
    "path": "LICENSE/LICENSE-FID",
    "content": "                                 Apache License\n                           Version 2.0, January 2004\n                        http://www.apache.org/licenses/\n\n   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION\n\n   1. Definitions.\n\n      \"License\" shall mean the terms and conditions for use, reproduction,\n      and distribution as defined by Sections 1 through 9 of this document.\n\n      \"Licensor\" shall mean the copyright owner or entity authorized by\n      the copyright owner that is granting the License.\n\n      \"Legal Entity\" shall mean the union of the acting entity and all\n      other entities that control, are controlled by, or are under common\n      control with that entity. For the purposes of this definition,\n      \"control\" means (i) the power, direct or indirect, to cause the\n      direction or management of such entity, whether by contract or\n      otherwise, or (ii) ownership of fifty percent (50%) or more of the\n      outstanding shares, or (iii) beneficial ownership of such entity.\n\n      \"You\" (or \"Your\") shall mean an individual or Legal Entity\n      exercising permissions granted by this License.\n\n      \"Source\" form shall mean the preferred form for making modifications,\n      including but not limited to software source code, documentation\n      source, and configuration files.\n\n      \"Object\" form shall mean any form resulting from mechanical\n      transformation or translation of a Source form, including but\n      not limited to compiled object code, generated documentation,\n      and conversions to other media types.\n\n      \"Work\" shall mean the work of authorship, whether in Source or\n      Object form, made available under the License, as indicated by a\n      copyright notice that is included in or attached to the work\n      (an example is provided in the Appendix below).\n\n      \"Derivative Works\" shall mean any work, whether in Source or Object\n      form, that is based on (or derived from) the Work and for which the\n      editorial revisions, annotations, elaborations, or other modifications\n      represent, as a whole, an original work of authorship. For the purposes\n      of this License, Derivative Works shall not include works that remain\n      separable from, or merely link (or bind by name) to the interfaces of,\n      the Work and Derivative Works thereof.\n\n      \"Contribution\" shall mean any work of authorship, including\n      the original version of the Work and any modifications or additions\n      to that Work or Derivative Works thereof, that is intentionally\n      submitted to Licensor for inclusion in the Work by the copyright owner\n      or by an individual or Legal Entity authorized to submit on behalf of\n      the copyright owner. For the purposes of this definition, \"submitted\"\n      means any form of electronic, verbal, or written communication sent\n      to the Licensor or its representatives, including but not limited to\n      communication on electronic mailing lists, source code control systems,\n      and issue tracking systems that are managed by, or on behalf of, the\n      Licensor for the purpose of discussing and improving the Work, but\n      excluding communication that is conspicuously marked or otherwise\n      designated in writing by the copyright owner as \"Not a Contribution.\"\n\n      \"Contributor\" shall mean Licensor and any individual or Legal Entity\n      on behalf of whom a Contribution has been received by Licensor and\n      subsequently incorporated within the Work.\n\n   2. Grant of Copyright License. Subject to the terms and conditions of\n      this License, each Contributor hereby grants to You a perpetual,\n      worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n      copyright license to reproduce, prepare Derivative Works of,\n      publicly display, publicly perform, sublicense, and distribute the\n      Work and such Derivative Works in Source or Object form.\n\n   3. Grant of Patent License. Subject to the terms and conditions of\n      this License, each Contributor hereby grants to You a perpetual,\n      worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n      (except as stated in this section) patent license to make, have made,\n      use, offer to sell, sell, import, and otherwise transfer the Work,\n      where such license applies only to those patent claims licensable\n      by such Contributor that are necessarily infringed by their\n      Contribution(s) alone or by combination of their Contribution(s)\n      with the Work to which such Contribution(s) was submitted. If You\n      institute patent litigation against any entity (including a\n      cross-claim or counterclaim in a lawsuit) alleging that the Work\n      or a Contribution incorporated within the Work constitutes direct\n      or contributory patent infringement, then any patent licenses\n      granted to You under this License for that Work shall terminate\n      as of the date such litigation is filed.\n\n   4. Redistribution. You may reproduce and distribute copies of the\n      Work or Derivative Works thereof in any medium, with or without\n      modifications, and in Source or Object form, provided that You\n      meet the following conditions:\n\n      (a) You must give any other recipients of the Work or\n          Derivative Works a copy of this License; and\n\n      (b) You must cause any modified files to carry prominent notices\n          stating that You changed the files; and\n\n      (c) You must retain, in the Source form of any Derivative Works\n          that You distribute, all copyright, patent, trademark, and\n          attribution notices from the Source form of the Work,\n          excluding those notices that do not pertain to any part of\n          the Derivative Works; and\n\n      (d) If the Work includes a \"NOTICE\" text file as part of its\n          distribution, then any Derivative Works that You distribute must\n          include a readable copy of the attribution notices contained\n          within such NOTICE file, excluding those notices that do not\n          pertain to any part of the Derivative Works, in at least one\n          of the following places: within a NOTICE text file distributed\n          as part of the Derivative Works; within the Source form or\n          documentation, if provided along with the Derivative Works; or,\n          within a display generated by the Derivative Works, if and\n          wherever such third-party notices normally appear. The contents\n          of the NOTICE file are for informational purposes only and\n          do not modify the License. You may add Your own attribution\n          notices within Derivative Works that You distribute, alongside\n          or as an addendum to the NOTICE text from the Work, provided\n          that such additional attribution notices cannot be construed\n          as modifying the License.\n\n      You may add Your own copyright statement to Your modifications and\n      may provide additional or different license terms and conditions\n      for use, reproduction, or distribution of Your modifications, or\n      for any such Derivative Works as a whole, provided Your use,\n      reproduction, and distribution of the Work otherwise complies with\n      the conditions stated in this License.\n\n   5. Submission of Contributions. Unless You explicitly state otherwise,\n      any Contribution intentionally submitted for inclusion in the Work\n      by You to the Licensor shall be under the terms and conditions of\n      this License, without any additional terms or conditions.\n      Notwithstanding the above, nothing herein shall supersede or modify\n      the terms of any separate license agreement you may have executed\n      with Licensor regarding such Contributions.\n\n   6. Trademarks. This License does not grant permission to use the trade\n      names, trademarks, service marks, or product names of the Licensor,\n      except as required for reasonable and customary use in describing the\n      origin of the Work and reproducing the content of the NOTICE file.\n\n   7. Disclaimer of Warranty. Unless required by applicable law or\n      agreed to in writing, Licensor provides the Work (and each\n      Contributor provides its Contributions) on an \"AS IS\" BASIS,\n      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or\n      implied, including, without limitation, any warranties or conditions\n      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A\n      PARTICULAR PURPOSE. You are solely responsible for determining the\n      appropriateness of using or redistributing the Work and assume any\n      risks associated with Your exercise of permissions under this License.\n\n   8. Limitation of Liability. In no event and under no legal theory,\n      whether in tort (including negligence), contract, or otherwise,\n      unless required by applicable law (such as deliberate and grossly\n      negligent acts) or agreed to in writing, shall any Contributor be\n      liable to You for damages, including any direct, indirect, special,\n      incidental, or consequential damages of any character arising as a\n      result of this License or out of the use or inability to use the\n      Work (including but not limited to damages for loss of goodwill,\n      work stoppage, computer failure or malfunction, or any and all\n      other commercial damages or losses), even if such Contributor\n      has been advised of the possibility of such damages.\n\n   9. Accepting Warranty or Additional Liability. While redistributing\n      the Work or Derivative Works thereof, You may choose to offer,\n      and charge a fee for, acceptance of support, warranty, indemnity,\n      or other liability obligations and/or rights consistent with this\n      License. However, in accepting such obligations, You may act only\n      on Your own behalf and on Your sole responsibility, not on behalf\n      of any other Contributor, and only if You agree to indemnify,\n      defend, and hold each Contributor harmless for any liability\n      incurred by, or claims asserted against, such Contributor by reason\n      of your accepting any such warranty or additional liability.\n\n   END OF TERMS AND CONDITIONS\n\n   APPENDIX: How to apply the Apache License to your work.\n\n      To apply the Apache License to your work, attach the following\n      boilerplate notice, with the fields enclosed by brackets \"[]\"\n      replaced with your own identifying information. (Don't include\n      the brackets!)  The text should be enclosed in the appropriate\n      comment syntax for the file format. We also recommend that a\n      file or class name and description of purpose be included on the\n      same \"printed page\" as the copyright notice for easier\n      identification within third-party archives.\n\n   Copyright [yyyy] [name of copyright owner]\n\n   Licensed under the Apache License, Version 2.0 (the \"License\");\n   you may not use this file except in compliance with the License.\n   You may obtain a copy of the License at\n\n       http://www.apache.org/licenses/LICENSE-2.0\n\n   Unless required by applicable law or agreed to in writing, software\n   distributed under the License is distributed on an \"AS IS\" BASIS,\n   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n   See the License for the specific language governing permissions and\n   limitations under the License.\n"
  },
  {
    "path": "LICENSE/LICENSE-LPIPS",
    "content": "Copyright (c) 2018, Richard Zhang, Phillip Isola, Alexei A. Efros, Eli Shechtman, Oliver Wang\r\nAll rights reserved.\r\n\r\nRedistribution and use in source and binary forms, with or without\r\nmodification, are permitted provided that the following conditions are met:\r\n\r\n* Redistributions of source code must retain the above copyright notice, this\r\n  list of conditions and the following disclaimer.\r\n\r\n* Redistributions in binary form must reproduce the above copyright notice,\r\n  this list of conditions and the following disclaimer in the documentation\r\n  and/or other materials provided with the distribution.\r\n\r\nTHIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\"\r\nAND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\r\nIMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\r\nDISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE\r\nFOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL\r\nDAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR\r\nSERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER\r\nCAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,\r\nOR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\r\nOF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\r\n\r\n"
  },
  {
    "path": "LICENSE/LICENSE-LUCIDRAINS",
    "content": "                    GNU GENERAL PUBLIC LICENSE\n                       Version 3, 29 June 2007\n\n Copyright (C) 2007 Free Software Foundation, Inc. <https://fsf.org/>\n Everyone is permitted to copy and distribute verbatim copies\n of this license document, but changing it is not allowed.\n\n                            Preamble\n\n  The GNU General Public License is a free, copyleft license for\nsoftware and other kinds of works.\n\n  The licenses for most software and other practical works are designed\nto take away your freedom to share and change the works.  By contrast,\nthe GNU General Public License is intended to guarantee your freedom to\nshare and change all versions of a program--to make sure it remains free\nsoftware for all its users.  We, the Free Software Foundation, use the\nGNU General Public License for most of our software; it applies also to\nany other work released this way by its authors.  You can apply it to\nyour programs, too.\n\n  When we speak of free software, we are referring to freedom, not\nprice.  Our General Public Licenses are designed to make sure that you\nhave the freedom to distribute copies of free software (and charge for\nthem if you wish), that you receive source code or can get it if you\nwant it, that you can change the software or use pieces of it in new\nfree programs, and that you know you can do these things.\n\n  To protect your rights, we need to prevent others from denying you\nthese rights or asking you to surrender the rights.  Therefore, you have\ncertain responsibilities if you distribute copies of the software, or if\nyou modify it: responsibilities to respect the freedom of others.\n\n  For example, if you distribute copies of such a program, whether\ngratis or for a fee, you must pass on to the recipients the same\nfreedoms that you received.  You must make sure that they, too, receive\nor can get the source code.  And you must show them these terms so they\nknow their rights.\n\n  Developers that use the GNU GPL protect your rights with two steps:\n(1) assert copyright on the software, and (2) offer you this License\ngiving you legal permission to copy, distribute and/or modify it.\n\n  For the developers' and authors' protection, the GPL clearly explains\nthat there is no warranty for this free software.  For both users' and\nauthors' sake, the GPL requires that modified versions be marked as\nchanged, so that their problems will not be attributed erroneously to\nauthors of previous versions.\n\n  Some devices are designed to deny users access to install or run\nmodified versions of the software inside them, although the manufacturer\ncan do so.  This is fundamentally incompatible with the aim of\nprotecting users' freedom to change the software.  The systematic\npattern of such abuse occurs in the area of products for individuals to\nuse, which is precisely where it is most unacceptable.  Therefore, we\nhave designed this version of the GPL to prohibit the practice for those\nproducts.  If such problems arise substantially in other domains, we\nstand ready to extend this provision to those domains in future versions\nof the GPL, as needed to protect the freedom of users.\n\n  Finally, every program is threatened constantly by software patents.\nStates should not allow patents to restrict development and use of\nsoftware on general-purpose computers, but in those that do, we wish to\navoid the special danger that patents applied to a free program could\nmake it effectively proprietary.  To prevent this, the GPL assures that\npatents cannot be used to render the program non-free.\n\n  The precise terms and conditions for copying, distribution and\nmodification follow.\n\n                       TERMS AND CONDITIONS\n\n  0. Definitions.\n\n  \"This License\" refers to version 3 of the GNU General Public License.\n\n  \"Copyright\" also means copyright-like laws that apply to other kinds of\nworks, such as semiconductor masks.\n\n  \"The Program\" refers to any copyrightable work licensed under this\nLicense.  Each licensee is addressed as \"you\".  \"Licensees\" and\n\"recipients\" may be individuals or organizations.\n\n  To \"modify\" a work means to copy from or adapt all or part of the work\nin a fashion requiring copyright permission, other than the making of an\nexact copy.  The resulting work is called a \"modified version\" of the\nearlier work or a work \"based on\" the earlier work.\n\n  A \"covered work\" means either the unmodified Program or a work based\non the Program.\n\n  To \"propagate\" a work means to do anything with it that, without\npermission, would make you directly or secondarily liable for\ninfringement under applicable copyright law, except executing it on a\ncomputer or modifying a private copy.  Propagation includes copying,\ndistribution (with or without modification), making available to the\npublic, and in some countries other activities as well.\n\n  To \"convey\" a work means any kind of propagation that enables other\nparties to make or receive copies.  Mere interaction with a user through\na computer network, with no transfer of a copy, is not conveying.\n\n  An interactive user interface displays \"Appropriate Legal Notices\"\nto the extent that it includes a convenient and prominently visible\nfeature that (1) displays an appropriate copyright notice, and (2)\ntells the user that there is no warranty for the work (except to the\nextent that warranties are provided), that licensees may convey the\nwork under this License, and how to view a copy of this License.  If\nthe interface presents a list of user commands or options, such as a\nmenu, a prominent item in the list meets this criterion.\n\n  1. Source Code.\n\n  The \"source code\" for a work means the preferred form of the work\nfor making modifications to it.  \"Object code\" means any non-source\nform of a work.\n\n  A \"Standard Interface\" means an interface that either is an official\nstandard defined by a recognized standards body, or, in the case of\ninterfaces specified for a particular programming language, one that\nis widely used among developers working in that language.\n\n  The \"System Libraries\" of an executable work include anything, other\nthan the work as a whole, that (a) is included in the normal form of\npackaging a Major Component, but which is not part of that Major\nComponent, and (b) serves only to enable use of the work with that\nMajor Component, or to implement a Standard Interface for which an\nimplementation is available to the public in source code form.  A\n\"Major Component\", in this context, means a major essential component\n(kernel, window system, and so on) of the specific operating system\n(if any) on which the executable work runs, or a compiler used to\nproduce the work, or an object code interpreter used to run it.\n\n  The \"Corresponding Source\" for a work in object code form means all\nthe source code needed to generate, install, and (for an executable\nwork) run the object code and to modify the work, including scripts to\ncontrol those activities.  However, it does not include the work's\nSystem Libraries, or general-purpose tools or generally available free\nprograms which are used unmodified in performing those activities but\nwhich are not part of the work.  For example, Corresponding Source\nincludes interface definition files associated with source files for\nthe work, and the source code for shared libraries and dynamically\nlinked subprograms that the work is specifically designed to require,\nsuch as by intimate data communication or control flow between those\nsubprograms and other parts of the work.\n\n  The Corresponding Source need not include anything that users\ncan regenerate automatically from other parts of the Corresponding\nSource.\n\n  The Corresponding Source for a work in source code form is that\nsame work.\n\n  2. Basic Permissions.\n\n  All rights granted under this License are granted for the term of\ncopyright on the Program, and are irrevocable provided the stated\nconditions are met.  This License explicitly affirms your unlimited\npermission to run the unmodified Program.  The output from running a\ncovered work is covered by this License only if the output, given its\ncontent, constitutes a covered work.  This License acknowledges your\nrights of fair use or other equivalent, as provided by copyright law.\n\n  You may make, run and propagate covered works that you do not\nconvey, without conditions so long as your license otherwise remains\nin force.  You may convey covered works to others for the sole purpose\nof having them make modifications exclusively for you, or provide you\nwith facilities for running those works, provided that you comply with\nthe terms of this License in conveying all material for which you do\nnot control copyright.  Those thus making or running the covered works\nfor you must do so exclusively on your behalf, under your direction\nand control, on terms that prohibit them from making any copies of\nyour copyrighted material outside their relationship with you.\n\n  Conveying under any other circumstances is permitted solely under\nthe conditions stated below.  Sublicensing is not allowed; section 10\nmakes it unnecessary.\n\n  3. Protecting Users' Legal Rights From Anti-Circumvention Law.\n\n  No covered work shall be deemed part of an effective technological\nmeasure under any applicable law fulfilling obligations under article\n11 of the WIPO copyright treaty adopted on 20 December 1996, or\nsimilar laws prohibiting or restricting circumvention of such\nmeasures.\n\n  When you convey a covered work, you waive any legal power to forbid\ncircumvention of technological measures to the extent such circumvention\nis effected by exercising rights under this License with respect to\nthe covered work, and you disclaim any intention to limit operation or\nmodification of the work as a means of enforcing, against the work's\nusers, your or third parties' legal rights to forbid circumvention of\ntechnological measures.\n\n  4. Conveying Verbatim Copies.\n\n  You may convey verbatim copies of the Program's source code as you\nreceive it, in any medium, provided that you conspicuously and\nappropriately publish on each copy an appropriate copyright notice;\nkeep intact all notices stating that this License and any\nnon-permissive terms added in accord with section 7 apply to the code;\nkeep intact all notices of the absence of any warranty; and give all\nrecipients a copy of this License along with the Program.\n\n  You may charge any price or no price for each copy that you convey,\nand you may offer support or warranty protection for a fee.\n\n  5. Conveying Modified Source Versions.\n\n  You may convey a work based on the Program, or the modifications to\nproduce it from the Program, in the form of source code under the\nterms of section 4, provided that you also meet all of these conditions:\n\n    a) The work must carry prominent notices stating that you modified\n    it, and giving a relevant date.\n\n    b) The work must carry prominent notices stating that it is\n    released under this License and any conditions added under section\n    7.  This requirement modifies the requirement in section 4 to\n    \"keep intact all notices\".\n\n    c) You must license the entire work, as a whole, under this\n    License to anyone who comes into possession of a copy.  This\n    License will therefore apply, along with any applicable section 7\n    additional terms, to the whole of the work, and all its parts,\n    regardless of how they are packaged.  This License gives no\n    permission to license the work in any other way, but it does not\n    invalidate such permission if you have separately received it.\n\n    d) If the work has interactive user interfaces, each must display\n    Appropriate Legal Notices; however, if the Program has interactive\n    interfaces that do not display Appropriate Legal Notices, your\n    work need not make them do so.\n\n  A compilation of a covered work with other separate and independent\nworks, which are not by their nature extensions of the covered work,\nand which are not combined with it such as to form a larger program,\nin or on a volume of a storage or distribution medium, is called an\n\"aggregate\" if the compilation and its resulting copyright are not\nused to limit the access or legal rights of the compilation's users\nbeyond what the individual works permit.  Inclusion of a covered work\nin an aggregate does not cause this License to apply to the other\nparts of the aggregate.\n\n  6. Conveying Non-Source Forms.\n\n  You may convey a covered work in object code form under the terms\nof sections 4 and 5, provided that you also convey the\nmachine-readable Corresponding Source under the terms of this License,\nin one of these ways:\n\n    a) Convey the object code in, or embodied in, a physical product\n    (including a physical distribution medium), accompanied by the\n    Corresponding Source fixed on a durable physical medium\n    customarily used for software interchange.\n\n    b) Convey the object code in, or embodied in, a physical product\n    (including a physical distribution medium), accompanied by a\n    written offer, valid for at least three years and valid for as\n    long as you offer spare parts or customer support for that product\n    model, to give anyone who possesses the object code either (1) a\n    copy of the Corresponding Source for all the software in the\n    product that is covered by this License, on a durable physical\n    medium customarily used for software interchange, for a price no\n    more than your reasonable cost of physically performing this\n    conveying of source, or (2) access to copy the\n    Corresponding Source from a network server at no charge.\n\n    c) Convey individual copies of the object code with a copy of the\n    written offer to provide the Corresponding Source.  This\n    alternative is allowed only occasionally and noncommercially, and\n    only if you received the object code with such an offer, in accord\n    with subsection 6b.\n\n    d) Convey the object code by offering access from a designated\n    place (gratis or for a charge), and offer equivalent access to the\n    Corresponding Source in the same way through the same place at no\n    further charge.  You need not require recipients to copy the\n    Corresponding Source along with the object code.  If the place to\n    copy the object code is a network server, the Corresponding Source\n    may be on a different server (operated by you or a third party)\n    that supports equivalent copying facilities, provided you maintain\n    clear directions next to the object code saying where to find the\n    Corresponding Source.  Regardless of what server hosts the\n    Corresponding Source, you remain obligated to ensure that it is\n    available for as long as needed to satisfy these requirements.\n\n    e) Convey the object code using peer-to-peer transmission, provided\n    you inform other peers where the object code and Corresponding\n    Source of the work are being offered to the general public at no\n    charge under subsection 6d.\n\n  A separable portion of the object code, whose source code is excluded\nfrom the Corresponding Source as a System Library, need not be\nincluded in conveying the object code work.\n\n  A \"User Product\" is either (1) a \"consumer product\", which means any\ntangible personal property which is normally used for personal, family,\nor household purposes, or (2) anything designed or sold for incorporation\ninto a dwelling.  In determining whether a product is a consumer product,\ndoubtful cases shall be resolved in favor of coverage.  For a particular\nproduct received by a particular user, \"normally used\" refers to a\ntypical or common use of that class of product, regardless of the status\nof the particular user or of the way in which the particular user\nactually uses, or expects or is expected to use, the product.  A product\nis a consumer product regardless of whether the product has substantial\ncommercial, industrial or non-consumer uses, unless such uses represent\nthe only significant mode of use of the product.\n\n  \"Installation Information\" for a User Product means any methods,\nprocedures, authorization keys, or other information required to install\nand execute modified versions of a covered work in that User Product from\na modified version of its Corresponding Source.  The information must\nsuffice to ensure that the continued functioning of the modified object\ncode is in no case prevented or interfered with solely because\nmodification has been made.\n\n  If you convey an object code work under this section in, or with, or\nspecifically for use in, a User Product, and the conveying occurs as\npart of a transaction in which the right of possession and use of the\nUser Product is transferred to the recipient in perpetuity or for a\nfixed term (regardless of how the transaction is characterized), the\nCorresponding Source conveyed under this section must be accompanied\nby the Installation Information.  But this requirement does not apply\nif neither you nor any third party retains the ability to install\nmodified object code on the User Product (for example, the work has\nbeen installed in ROM).\n\n  The requirement to provide Installation Information does not include a\nrequirement to continue to provide support service, warranty, or updates\nfor a work that has been modified or installed by the recipient, or for\nthe User Product in which it has been modified or installed.  Access to a\nnetwork may be denied when the modification itself materially and\nadversely affects the operation of the network or violates the rules and\nprotocols for communication across the network.\n\n  Corresponding Source conveyed, and Installation Information provided,\nin accord with this section must be in a format that is publicly\ndocumented (and with an implementation available to the public in\nsource code form), and must require no special password or key for\nunpacking, reading or copying.\n\n  7. Additional Terms.\n\n  \"Additional permissions\" are terms that supplement the terms of this\nLicense by making exceptions from one or more of its conditions.\nAdditional permissions that are applicable to the entire Program shall\nbe treated as though they were included in this License, to the extent\nthat they are valid under applicable law.  If additional permissions\napply only to part of the Program, that part may be used separately\nunder those permissions, but the entire Program remains governed by\nthis License without regard to the additional permissions.\n\n  When you convey a copy of a covered work, you may at your option\nremove any additional permissions from that copy, or from any part of\nit.  (Additional permissions may be written to require their own\nremoval in certain cases when you modify the work.)  You may place\nadditional permissions on material, added by you to a covered work,\nfor which you have or can give appropriate copyright permission.\n\n  Notwithstanding any other provision of this License, for material you\nadd to a covered work, you may (if authorized by the copyright holders of\nthat material) supplement the terms of this License with terms:\n\n    a) Disclaiming warranty or limiting liability differently from the\n    terms of sections 15 and 16 of this License; or\n\n    b) Requiring preservation of specified reasonable legal notices or\n    author attributions in that material or in the Appropriate Legal\n    Notices displayed by works containing it; or\n\n    c) Prohibiting misrepresentation of the origin of that material, or\n    requiring that modified versions of such material be marked in\n    reasonable ways as different from the original version; or\n\n    d) Limiting the use for publicity purposes of names of licensors or\n    authors of the material; or\n\n    e) Declining to grant rights under trademark law for use of some\n    trade names, trademarks, or service marks; or\n\n    f) Requiring indemnification of licensors and authors of that\n    material by anyone who conveys the material (or modified versions of\n    it) with contractual assumptions of liability to the recipient, for\n    any liability that these contractual assumptions directly impose on\n    those licensors and authors.\n\n  All other non-permissive additional terms are considered \"further\nrestrictions\" within the meaning of section 10.  If the Program as you\nreceived it, or any part of it, contains a notice stating that it is\ngoverned by this License along with a term that is a further\nrestriction, you may remove that term.  If a license document contains\na further restriction but permits relicensing or conveying under this\nLicense, you may add to a covered work material governed by the terms\nof that license document, provided that the further restriction does\nnot survive such relicensing or conveying.\n\n  If you add terms to a covered work in accord with this section, you\nmust place, in the relevant source files, a statement of the\nadditional terms that apply to those files, or a notice indicating\nwhere to find the applicable terms.\n\n  Additional terms, permissive or non-permissive, may be stated in the\nform of a separately written license, or stated as exceptions;\nthe above requirements apply either way.\n\n  8. Termination.\n\n  You may not propagate or modify a covered work except as expressly\nprovided under this License.  Any attempt otherwise to propagate or\nmodify it is void, and will automatically terminate your rights under\nthis License (including any patent licenses granted under the third\nparagraph of section 11).\n\n  However, if you cease all violation of this License, then your\nlicense from a particular copyright holder is reinstated (a)\nprovisionally, unless and until the copyright holder explicitly and\nfinally terminates your license, and (b) permanently, if the copyright\nholder fails to notify you of the violation by some reasonable means\nprior to 60 days after the cessation.\n\n  Moreover, your license from a particular copyright holder is\nreinstated permanently if the copyright holder notifies you of the\nviolation by some reasonable means, this is the first time you have\nreceived notice of violation of this License (for any work) from that\ncopyright holder, and you cure the violation prior to 30 days after\nyour receipt of the notice.\n\n  Termination of your rights under this section does not terminate the\nlicenses of parties who have received copies or rights from you under\nthis License.  If your rights have been terminated and not permanently\nreinstated, you do not qualify to receive new licenses for the same\nmaterial under section 10.\n\n  9. Acceptance Not Required for Having Copies.\n\n  You are not required to accept this License in order to receive or\nrun a copy of the Program.  Ancillary propagation of a covered work\noccurring solely as a consequence of using peer-to-peer transmission\nto receive a copy likewise does not require acceptance.  However,\nnothing other than this License grants you permission to propagate or\nmodify any covered work.  These actions infringe copyright if you do\nnot accept this License.  Therefore, by modifying or propagating a\ncovered work, you indicate your acceptance of this License to do so.\n\n  10. Automatic Licensing of Downstream Recipients.\n\n  Each time you convey a covered work, the recipient automatically\nreceives a license from the original licensors, to run, modify and\npropagate that work, subject to this License.  You are not responsible\nfor enforcing compliance by third parties with this License.\n\n  An \"entity transaction\" is a transaction transferring control of an\norganization, or substantially all assets of one, or subdividing an\norganization, or merging organizations.  If propagation of a covered\nwork results from an entity transaction, each party to that\ntransaction who receives a copy of the work also receives whatever\nlicenses to the work the party's predecessor in interest had or could\ngive under the previous paragraph, plus a right to possession of the\nCorresponding Source of the work from the predecessor in interest, if\nthe predecessor has it or can get it with reasonable efforts.\n\n  You may not impose any further restrictions on the exercise of the\nrights granted or affirmed under this License.  For example, you may\nnot impose a license fee, royalty, or other charge for exercise of\nrights granted under this License, and you may not initiate litigation\n(including a cross-claim or counterclaim in a lawsuit) alleging that\nany patent claim is infringed by making, using, selling, offering for\nsale, or importing the Program or any portion of it.\n\n  11. Patents.\n\n  A \"contributor\" is a copyright holder who authorizes use under this\nLicense of the Program or a work on which the Program is based.  The\nwork thus licensed is called the contributor's \"contributor version\".\n\n  A contributor's \"essential patent claims\" are all patent claims\nowned or controlled by the contributor, whether already acquired or\nhereafter acquired, that would be infringed by some manner, permitted\nby this License, of making, using, or selling its contributor version,\nbut do not include claims that would be infringed only as a\nconsequence of further modification of the contributor version.  For\npurposes of this definition, \"control\" includes the right to grant\npatent sublicenses in a manner consistent with the requirements of\nthis License.\n\n  Each contributor grants you a non-exclusive, worldwide, royalty-free\npatent license under the contributor's essential patent claims, to\nmake, use, sell, offer for sale, import and otherwise run, modify and\npropagate the contents of its contributor version.\n\n  In the following three paragraphs, a \"patent license\" is any express\nagreement or commitment, however denominated, not to enforce a patent\n(such as an express permission to practice a patent or covenant not to\nsue for patent infringement).  To \"grant\" such a patent license to a\nparty means to make such an agreement or commitment not to enforce a\npatent against the party.\n\n  If you convey a covered work, knowingly relying on a patent license,\nand the Corresponding Source of the work is not available for anyone\nto copy, free of charge and under the terms of this License, through a\npublicly available network server or other readily accessible means,\nthen you must either (1) cause the Corresponding Source to be so\navailable, or (2) arrange to deprive yourself of the benefit of the\npatent license for this particular work, or (3) arrange, in a manner\nconsistent with the requirements of this License, to extend the patent\nlicense to downstream recipients.  \"Knowingly relying\" means you have\nactual knowledge that, but for the patent license, your conveying the\ncovered work in a country, or your recipient's use of the covered work\nin a country, would infringe one or more identifiable patents in that\ncountry that you have reason to believe are valid.\n\n  If, pursuant to or in connection with a single transaction or\narrangement, you convey, or propagate by procuring conveyance of, a\ncovered work, and grant a patent license to some of the parties\nreceiving the covered work authorizing them to use, propagate, modify\nor convey a specific copy of the covered work, then the patent license\nyou grant is automatically extended to all recipients of the covered\nwork and works based on it.\n\n  A patent license is \"discriminatory\" if it does not include within\nthe scope of its coverage, prohibits the exercise of, or is\nconditioned on the non-exercise of one or more of the rights that are\nspecifically granted under this License.  You may not convey a covered\nwork if you are a party to an arrangement with a third party that is\nin the business of distributing software, under which you make payment\nto the third party based on the extent of your activity of conveying\nthe work, and under which the third party grants, to any of the\nparties who would receive the covered work from you, a discriminatory\npatent license (a) in connection with copies of the covered work\nconveyed by you (or copies made from those copies), or (b) primarily\nfor and in connection with specific products or compilations that\ncontain the covered work, unless you entered into that arrangement,\nor that patent license was granted, prior to 28 March 2007.\n\n  Nothing in this License shall be construed as excluding or limiting\nany implied license or other defenses to infringement that may\notherwise be available to you under applicable patent law.\n\n  12. No Surrender of Others' Freedom.\n\n  If conditions are imposed on you (whether by court order, agreement or\notherwise) that contradict the conditions of this License, they do not\nexcuse you from the conditions of this License.  If you cannot convey a\ncovered work so as to satisfy simultaneously your obligations under this\nLicense and any other pertinent obligations, then as a consequence you may\nnot convey it at all.  For example, if you agree to terms that obligate you\nto collect a royalty for further conveying from those to whom you convey\nthe Program, the only way you could satisfy both those terms and this\nLicense would be to refrain entirely from conveying the Program.\n\n  13. Use with the GNU Affero General Public License.\n\n  Notwithstanding any other provision of this License, you have\npermission to link or combine any covered work with a work licensed\nunder version 3 of the GNU Affero General Public License into a single\ncombined work, and to convey the resulting work.  The terms of this\nLicense will continue to apply to the part which is the covered work,\nbut the special requirements of the GNU Affero General Public License,\nsection 13, concerning interaction through a network will apply to the\ncombination as such.\n\n  14. Revised Versions of this License.\n\n  The Free Software Foundation may publish revised and/or new versions of\nthe GNU General Public License from time to time.  Such new versions will\nbe similar in spirit to the present version, but may differ in detail to\naddress new problems or concerns.\n\n  Each version is given a distinguishing version number.  If the\nProgram specifies that a certain numbered version of the GNU General\nPublic License \"or any later version\" applies to it, you have the\noption of following the terms and conditions either of that numbered\nversion or of any later version published by the Free Software\nFoundation.  If the Program does not specify a version number of the\nGNU General Public License, you may choose any version ever published\nby the Free Software Foundation.\n\n  If the Program specifies that a proxy can decide which future\nversions of the GNU General Public License can be used, that proxy's\npublic statement of acceptance of a version permanently authorizes you\nto choose that version for the Program.\n\n  Later license versions may give you additional or different\npermissions.  However, no additional obligations are imposed on any\nauthor or copyright holder as a result of your choosing to follow a\nlater version.\n\n  15. Disclaimer of Warranty.\n\n  THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY\nAPPLICABLE LAW.  EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT\nHOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM \"AS IS\" WITHOUT WARRANTY\nOF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO,\nTHE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR\nPURPOSE.  THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM\nIS WITH YOU.  SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF\nALL NECESSARY SERVICING, REPAIR OR CORRECTION.\n\n  16. Limitation of Liability.\n\n  IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING\nWILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS\nTHE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY\nGENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE\nUSE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF\nDATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD\nPARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS),\nEVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF\nSUCH DAMAGES.\n\n  17. Interpretation of Sections 15 and 16.\n\n  If the disclaimer of warranty and limitation of liability provided\nabove cannot be given local legal effect according to their terms,\nreviewing courts shall apply local law that most closely approximates\nan absolute waiver of all civil liability in connection with the\nProgram, unless a warranty or assumption of liability accompanies a\ncopy of the Program in return for a fee.\n\n                     END OF TERMS AND CONDITIONS\n\n            How to Apply These Terms to Your New Programs\n\n  If you develop a new program, and you want it to be of the greatest\npossible use to the public, the best way to achieve this is to make it\nfree software which everyone can redistribute and change under these terms.\n\n  To do so, attach the following notices to the program.  It is safest\nto attach them to the start of each source file to most effectively\nstate the exclusion of warranty; and each file should have at least\nthe \"copyright\" line and a pointer to where the full notice is found.\n\n    <one line to give the program's name and a brief idea of what it does.>\n    Copyright (C) <year>  <name of author>\n\n    This program is free software: you can redistribute it and/or modify\n    it under the terms of the GNU General Public License as published by\n    the Free Software Foundation, either version 3 of the License, or\n    (at your option) any later version.\n\n    This program is distributed in the hope that it will be useful,\n    but WITHOUT ANY WARRANTY; without even the implied warranty of\n    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the\n    GNU General Public License for more details.\n\n    You should have received a copy of the GNU General Public License\n    along with this program.  If not, see <https://www.gnu.org/licenses/>.\n\nAlso add information on how to contact you by electronic and paper mail.\n\n  If the program does terminal interaction, make it output a short\nnotice like this when it starts in an interactive mode:\n\n    <program>  Copyright (C) <year>  <name of author>\n    This program comes with ABSOLUTELY NO WARRANTY; for details type `show w'.\n    This is free software, and you are welcome to redistribute it\n    under certain conditions; type `show c' for details.\n\nThe hypothetical commands `show w' and `show c' should show the appropriate\nparts of the General Public License.  Of course, your program's commands\nmight be different; for a GUI interface, you would use an \"about box\".\n\n  You should also get your employer (if you work as a programmer) or school,\nif any, to sign a \"copyright disclaimer\" for the program, if necessary.\nFor more information on this, and how to apply and follow the GNU GPL, see\n<https://www.gnu.org/licenses/>.\n\n  The GNU General Public License does not permit incorporating your program\ninto proprietary programs.  If your program is a subroutine library, you\nmay consider it more useful to permit linking proprietary applications with\nthe library.  If this is what you want to do, use the GNU Lesser General\nPublic License instead of this License.  But first, please read\n<https://www.gnu.org/licenses/why-not-lgpl.html>.\n"
  },
  {
    "path": "LICENSE/LICENSE-NVIDIA",
    "content": "Copyright (c) 2019, NVIDIA Corporation. All rights reserved.\r\n\r\n\r\nNvidia Source Code License-NC\r\n\r\n=======================================================================\r\n\r\n1. Definitions\r\n\r\n\"Licensor\" means any person or entity that distributes its Work.\r\n\r\n\"Software\" means the original work of authorship made available under\r\nthis License.\r\n\r\n\"Work\" means the Software and any additions to or derivative works of\r\nthe Software that are made available under this License.\r\n\r\n\"Nvidia Processors\" means any central processing unit (CPU), graphics\r\nprocessing unit (GPU), field-programmable gate array (FPGA),\r\napplication-specific integrated circuit (ASIC) or any combination\r\nthereof designed, made, sold, or provided by Nvidia or its affiliates.\r\n\r\nThe terms \"reproduce,\" \"reproduction,\" \"derivative works,\" and\r\n\"distribution\" have the meaning as provided under U.S. copyright law;\r\nprovided, however, that for the purposes of this License, derivative\r\nworks shall not include works that remain separable from, or merely\r\nlink (or bind by name) to the interfaces of, the Work.\r\n\r\nWorks, including the Software, are \"made available\" under this License\r\nby including in or with the Work either (a) a copyright notice\r\nreferencing the applicability of this License to the Work, or (b) a\r\ncopy of this License.\r\n\r\n2. License Grants\r\n\r\n    2.1 Copyright Grant. Subject to the terms and conditions of this\r\n    License, each Licensor grants to you a perpetual, worldwide,\r\n    non-exclusive, royalty-free, copyright license to reproduce,\r\n    prepare derivative works of, publicly display, publicly perform,\r\n    sublicense and distribute its Work and any resulting derivative\r\n    works in any form.\r\n\r\n3. Limitations\r\n\r\n    3.1 Redistribution. You may reproduce or distribute the Work only\r\n    if (a) you do so under this License, (b) you include a complete\r\n    copy of this License with your distribution, and (c) you retain\r\n    without modification any copyright, patent, trademark, or\r\n    attribution notices that are present in the Work.\r\n\r\n    3.2 Derivative Works. You may specify that additional or different\r\n    terms apply to the use, reproduction, and distribution of your\r\n    derivative works of the Work (\"Your Terms\") only if (a) Your Terms\r\n    provide that the use limitation in Section 3.3 applies to your\r\n    derivative works, and (b) you identify the specific derivative\r\n    works that are subject to Your Terms. Notwithstanding Your Terms,\r\n    this License (including the redistribution requirements in Section\r\n    3.1) will continue to apply to the Work itself.\r\n\r\n    3.3 Use Limitation. The Work and any derivative works thereof only\r\n    may be used or intended for use non-commercially. The Work or\r\n    derivative works thereof may be used or intended for use by Nvidia\r\n    or its affiliates commercially or non-commercially. As used herein,\r\n    \"non-commercially\" means for research or evaluation purposes only.\r\n\r\n    3.4 Patent Claims. If you bring or threaten to bring a patent claim\r\n    against any Licensor (including any claim, cross-claim or\r\n    counterclaim in a lawsuit) to enforce any patents that you allege\r\n    are infringed by any Work, then your rights under this License from\r\n    such Licensor (including the grants in Sections 2.1 and 2.2) will\r\n    terminate immediately.\r\n\r\n    3.5 Trademarks. This License does not grant any rights to use any\r\n    Licensor's or its affiliates' names, logos, or trademarks, except\r\n    as necessary to reproduce the notices described in this License.\r\n\r\n    3.6 Termination. If you violate any term of this License, then your\r\n    rights under this License (including the grants in Sections 2.1 and\r\n    2.2) will terminate immediately.\r\n\r\n4. Disclaimer of Warranty.\r\n\r\nTHE WORK IS PROVIDED \"AS IS\" WITHOUT WARRANTIES OR CONDITIONS OF ANY\r\nKIND, EITHER EXPRESS OR IMPLIED, INCLUDING WARRANTIES OR CONDITIONS OF\r\nMERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, TITLE OR\r\nNON-INFRINGEMENT. YOU BEAR THE RISK OF UNDERTAKING ANY ACTIVITIES UNDER\r\nTHIS LICENSE. \r\n\r\n5. Limitation of Liability.\r\n\r\nEXCEPT AS PROHIBITED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL\r\nTHEORY, WHETHER IN TORT (INCLUDING NEGLIGENCE), CONTRACT, OR OTHERWISE\r\nSHALL ANY LICENSOR BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY DIRECT,\r\nINDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES ARISING OUT OF\r\nOR RELATED TO THIS LICENSE, THE USE OR INABILITY TO USE THE WORK\r\n(INCLUDING BUT NOT LIMITED TO LOSS OF GOODWILL, BUSINESS INTERRUPTION,\r\nLOST PROFITS OR DATA, COMPUTER FAILURE OR MALFUNCTION, OR ANY OTHER\r\nCOMMERCIAL DAMAGES OR LOSSES), EVEN IF THE LICENSOR HAS BEEN ADVISED OF\r\nTHE POSSIBILITY OF SUCH DAMAGES.\r\n\r\n=======================================================================\r\n"
  },
  {
    "path": "LICENSE/LICENSE-ROSINALITY",
    "content": "MIT License\n\nCopyright (c) 2019 Kim Seonghyeon\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n"
  },
  {
    "path": "LICENSE/LICENSE-VGG",
    "content": "Copyright (C) 2019 NVIDIA Corporation. Ting-Chun Wang, Ming-Yu Liu, Jun-Yan Zhu.\nBSD License. All rights reserved. \n\nRedistribution and use in source and binary forms, with or without\nmodification, are permitted provided that the following conditions are met:\n\n* Redistributions of source code must retain the above copyright notice, this\n  list of conditions and the following disclaimer.\n\n* Redistributions in binary form must reproduce the above copyright notice,\n  this list of conditions and the following disclaimer in the documentation\n  and/or other materials provided with the distribution.\n\nTHE AUTHOR DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS SOFTWARE, INCLUDING ALL \nIMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR ANY PARTICULAR PURPOSE. \nIN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, INDIRECT OR CONSEQUENTIAL \nDAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, \nWHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING \nOUT OF OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.\n\n\n--------------------------- LICENSE FOR pytorch-CycleGAN-and-pix2pix ----------------\nCopyright (c) 2017, Jun-Yan Zhu and Taesung Park\nAll rights reserved.\n\nRedistribution and use in source and binary forms, with or without\nmodification, are permitted provided that the following conditions are met:\n\n* Redistributions of source code must retain the above copyright notice, this\n  list of conditions and the following disclaimer.\n\n* Redistributions in binary form must reproduce the above copyright notice,\n  this list of conditions and the following disclaimer in the documentation\n  and/or other materials provided with the distribution.\n\nTHIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\"\nAND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\nIMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\nDISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE\nFOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL\nDAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR\nSERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER\nCAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,\nOR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\nOF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE."
  },
  {
    "path": "README.md",
    "content": "# maua-stylegan2\r\n\r\nThis is the repo for my experiments with StyleGAN2. There are many like it, but this one is mine.\r\n\r\nIt 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/).\r\n\r\nThe 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.\r\n\r\nThe files/folders of interest and their purpose are:\r\n\r\n| File/Folder | Description\r\n| :--- | :----------\r\n| generate_audiovisual.py | used to generate audio-reactive interpolations\r\n| audioreactive/ | contains the main functions needed for audioreactiveness + examples demonstrating how they can be used\r\n| render.py | renders interpolations using ffmpeg\r\n| select_latents.py | GUI for selecting latents, left click to add to top set, right click to add to bottom\r\n| models/ | StyleGAN networks\r\n| workspace/ | place to store intermediate results, latents, or inputs, etc.\r\n| output/ | default generated output folder\r\n| train.py | code for training models\r\n\r\nThe rest of the code is experimental, probably broken, and unsupported.\r\n\r\n## Installation\r\n\r\n```bash\r\ngit clone https://github.com/JCBrouwer/maua-stylegan2\r\ncd maua-stylegan2\r\npip install -r requirements.txt\r\n```\r\n\r\nAlternatively, check out this [Colab Notebook](https://colab.research.google.com/drive/1Ig1EXfmBC01qik11Q32P0ZffFtNipiBR)\r\n\r\n## Generating audio-reactive interpolations\r\n\r\nThe simplest way to get started is to try either (in shell):\r\n```bash\r\npython generate_audiovisual.py --ckpt \"/path/to/model.pt\" --audio_file \"/path/to/audio.wav\"\r\n```\r\nor (in e.g. a jupyter notebook):\r\n```python\r\nfrom generate_audiovisual import generate\r\ngenerate(\"/path/to/model.pt\", \"/path/to/audio.wav\")\r\n```\r\n\r\nThis will use the default audio-reactive settings (which aren't great).\r\n\r\nTo customize the generated interpolation, more functions can be defined to generate latents, noise, network bends, model rewrites, and truncation.\r\n\r\n```python\r\nimport audioreactive as ar\r\nfrom generate_audiovisual import generate\r\n\r\ndef initialize(args):\r\n    args.onsets = ar.onsets(args.audio, args.sr, ...)\r\n    args.chroma = ar.chroma(args.audio, args.sr, ...)\r\n    return args\r\n\r\ndef get_latents(selection, args):\r\n    latents = ar.chroma_weight_latents(args.chroma, selection)\r\n    return latents\r\n\r\ndef get_noise(height, width, scale, num_scales, args):\r\n    noise = ar.perlin_noise(...)\r\n    noise *= 1 + args.onsets\r\n    return noise\r\n\r\ngenerate(ckpt=\"/path/to/model.pt\", audio_file=\"/path/to/audio.wav\", initialize=initialize, get_latents=get_latents, get_noise=get_noise)\r\n```\r\n\r\nWhen running from command line, the `generate()` call at the end can be left out and the interpolation can be generated with:\r\n\r\n```bash\r\npython generate_audiovisual.py --ckpt \"/path/to/model.pt\" --audio_file \"/path/to/audio.wav\" --audioreactive_file \"/path/to/the/code_above.py\"\r\n```\r\n\r\nThis 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).\r\n\r\nWithin 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`).\r\n\r\nExamples 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).\r\n\r\nOne important thing to note is that the outputs of the functions must adhere strictly to the expected formats. \r\n\r\nEach 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:\r\n- audio: raw audio signal\r\n- sr: sampling rate of audio\r\n- n_frames: total number of interpolation frames\r\n- duration: length of audio in seconds\r\n\r\n```python\r\ndef initialize(args):\r\n    # intialize values used in multiple of the following functions here\r\n    # e.g. onsets, chroma, RMS, segmentations, bpms, etc.\r\n    # this is useful to prevent duplicate computations (get_noise is called for each noise size)\r\n    # remember to store them back in args\r\n    ...\r\n    return args\r\n\r\ndef get_latents(selection, args):\r\n    # selection holds some latent vectors (generated randomly or from a file)\r\n    # generate an audioreactive latent tensor of shape [n_frames, layers, latent_dim]\r\n    ...\r\n    return latents\r\n\r\ndef get_noise(height, width, scale, num_scales, args):\r\n    # height and width are the spatial dimensions of the current noise layer\r\n    # scale is the index and num_scales the total number of noise layers\r\n    # generate an audioreactive noise tensor of shape [n_frames, 1, height, width]\r\n    ...\r\n    return noise\r\n\r\ndef get_bends(args):\r\n    # generate a list of dictionaries specifying network bends\r\n    # these must follow one of two forms:\r\n    #\r\n    # either: {\r\n    #     \"layer\": layer index to apply bend to,\r\n    #     \"transform\": torch.nn.Module that applies the transformation,\r\n    # }\r\n    # or: {\r\n    #     \"layer\": layer index to apply bend to,\r\n    #     \"modulation\": time dependent modulation of the transformation, shape=(n_frames, ...), \r\n    #     \"transform\": function that takes a batch of modulation and returns a torch.nn.Module\r\n    #                  that applies the transformation (given the modulation batch),\r\n    # }\r\n    # (The second one is technical debt in a nutshell. It's a workaround to get kornia transforms\r\n    #  to play nicely. You're probably better off using the first option with a th.nn.Module that\r\n    #  has its modulation as an attribute and keeps count of which frame it's rendering internally).\r\n    ...\r\n    return bends\r\n\r\ndef get_rewrites(args):\r\n    # generate a dictionary specifying model rewrites\r\n    # each key value pair should follow:\r\n    #       param_name -> [transform, modulation]\r\n    # where: param_name is the fully-qualified parameter name (see generator.named_children())\r\n    #        transform & modulation follow the form of the second network bending dict option above\r\n    ...\r\n    return rewrites\r\n\r\ndef get_truncation(args):\r\n    # generate a sequence of truncation values of shape (n_frames,)\r\n    ...\r\n    return truncation\r\n```\r\n\r\nThe arguments to `generate_audiovisual.py` are as follows. The first two are required, and the remaining are optional.\r\n```bash\r\ngenerate_audiovisual.py\r\n  --ckpt CKPT                              # path to model checkpoint\r\n  --audio_file AUDIO_FILE                  # path to audio file to react to\r\n  --audioreactive_file AUDIOREACTIVE_FILE  # file with audio-reactive functions defined (as above)\r\n  --output_dir OUTPUT_DIR                  # path to output dir\r\n  --offset OFFSET                          # starting time in audio in seconds (defaults to 0)\r\n  --duration DURATION                      # duration of interpolation to generate in seconds (leave empty for length of audiofile)\r\n  --latent_file LATENT_FILE                # path to latents saved as numpy array\r\n  --shuffle_latents                        # whether to shuffle the supplied latents or not\r\n  --out_size OUT_SIZE                      # ouput video size: [512, 1024, or 1920]\r\n  --fps FPS                                # output video framerate\r\n  --batch BATCH                            # batch size to render with\r\n  --truncation TRUNCATION                  # truncation to render with (leave empty if get_truncations() is in --audioreactive_file)\r\n  --randomize_noise                        # whether to randomize noise\r\n  --dataparallel                           # whether to use data parallel rendering\r\n  --stylegan1                              # if the model checkpoint is StyleGAN1\r\n  --G_res G_RES                            # training resolution of the generator\r\n  --base_res_factor BASE_RES_FACTOR        # factor to increase generator noise maps by (useful when e.g. doubling 512px net to 1024px)\r\n  --noconst                                # whether the generator was trained without a constant input layer\r\n  --latent_dim LATENT_DIM                  # latent vector size of the generator\r\n  --n_mlp N_MLP                            # number of mapping network layers\r\n  --channel_multiplier CHANNEL_MULTIPLIER  # generator's channel scaling multiplier\r\n```\r\n\r\nAlternatively, `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).\r\n\r\nModel 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).\r\n\r\nThere is minimal support for rendering with StyleGAN1 checkpoints as well, although only with latent and noise (no network bending or model rewriting).\r\n\r\n## Citation\r\n\r\nIf you use the techniques introduced in the paper or the code in this repository for your research, please cite the paper:\r\n```\r\n@InProceedings{Brouwer_2020_NeurIPS_Workshops},\r\n    author = {Brouwer, Hans},\r\n    title = {Audio-reactive Latent Interpolations with StyleGAN},\r\n    booktitle = {Proceedings of the 4th Workshop on Machine Learning for Creativity and Design at NeurIPS 2020},\r\n    month = {December},\r\n    year = {2020},\r\n    url={https://jcbrouwer.github.io/assets/audio-reactive-stylegan/paper.pdf}\r\n}\r\n```\r\n"
  },
  {
    "path": "accelerate/accelerate_inception.py",
    "content": "import os\nimport gc\nimport wandb\nimport argparse\nimport torch as th\nfrom tqdm import tqdm\nfrom torch.utils import data\nimport torch.nn.functional as F\nfrom inception_vae import InceptionVAE\nfrom dataset import MultiResolutionDataset\nfrom torchvision import transforms, utils, models\n\n\ndef info(x):\n    print(x.shape, x.detach().cpu().min(), x.detach().cpu().mean(), x.detach().cpu().max())\n\n\ndef sample_data(loader):\n    while True:\n        for batch in loader:\n            yield batch\n\n\nclass VGG19(th.nn.Module):\n    \"\"\"\n    Adapted from https://github.com/NVIDIA/pix2pixHD\n    See LICENSE-VGG\n    \"\"\"\n\n    def __init__(self, requires_grad=False):\n        super(VGG19, self).__init__()\n        vgg_pretrained_features = models.vgg19(pretrained=True).features\n        self.slice1 = th.nn.Sequential()\n        self.slice2 = th.nn.Sequential()\n        self.slice3 = th.nn.Sequential()\n        self.slice4 = th.nn.Sequential()\n        self.slice5 = th.nn.Sequential()\n        for x in range(2):\n            self.slice1.add_module(str(x), vgg_pretrained_features[x])\n        for x in range(2, 7):\n            self.slice2.add_module(str(x), vgg_pretrained_features[x])\n        for x in range(7, 12):\n            self.slice3.add_module(str(x), vgg_pretrained_features[x])\n        for x in range(12, 21):\n            self.slice4.add_module(str(x), vgg_pretrained_features[x])\n        for x in range(21, 30):\n            self.slice5.add_module(str(x), vgg_pretrained_features[x])\n        if not requires_grad:\n            for param in self.parameters():\n                param.requires_grad = False\n\n    def forward(self, X):\n        h_relu1 = self.slice1(X)\n        h_relu2 = self.slice2(h_relu1)\n        h_relu3 = self.slice3(h_relu2)\n        h_relu4 = self.slice4(h_relu3)\n        h_relu5 = self.slice5(h_relu4)\n        out = [h_relu1, h_relu2, h_relu3, h_relu4, h_relu5]\n        return out\n\n\nclass VGGLoss(th.nn.Module):\n    \"\"\"\n    Adapted from https://github.com/NVIDIA/pix2pixHD\n    See LICENSE-VGG\n    \"\"\"\n\n    def __init__(self):\n        super(VGGLoss, self).__init__()\n        self.vgg = VGG19().cuda()\n        self.criterion = th.nn.L1Loss()\n        self.weights = [1.0 / 32, 1.0 / 16, 1.0 / 8, 1.0 / 4, 1.0]\n\n    def forward(self, x, y):\n        x_vgg, y_vgg = self.vgg(x), self.vgg(y)\n        loss = 0\n        for i in range(len(x_vgg)):\n            loss += self.weights[i] * self.criterion(x_vgg[i], y_vgg[i].detach())\n        return loss\n\n\ndef train(latent_dim, num_repeats, learning_rate, lambda_vgg, lambda_mse):\n    print(\n        f\"latent_dim={latent_dim:.4f}\",\n        f\"num_repeats={num_repeats:.4f}\",\n        f\"learning_rate={learning_rate:.4f}\",\n        f\"lambda_vgg={lambda_vgg:.4f}\",\n        f\"lambda_mse={lambda_mse:.4f}\",\n    )\n\n    transform = transforms.Compose(\n        [\n            transforms.Resize(128),\n            transforms.RandomHorizontalFlip(p=0.5),\n            transforms.ToTensor(),\n            # transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True),\n        ]\n    )\n    batch_size = 72\n    data_path = \"/home/hans/trainsets/cyphis\"\n    name = os.path.splitext(os.path.basename(data_path))[0]\n    dataset = MultiResolutionDataset(data_path, transform, 256)\n    dataloader = data.DataLoader(\n        dataset, batch_size=batch_size, sampler=data.RandomSampler(dataset), num_workers=12, drop_last=True,\n    )\n    loader = sample_data(dataloader)\n    sample_imgs = next(loader)[:24]\n    wandb.log({\"Real Images\": [wandb.Image(utils.make_grid(sample_imgs, nrow=6, normalize=True, range=(0, 1)))]})\n\n    vae, vae_optim = None, None\n    vae = InceptionVAE(latent_dim=latent_dim, repeat_per_block=num_repeats).to(device)\n    vae_optim = th.optim.Adam(vae.parameters(), lr=learning_rate)\n\n    vgg = VGGLoss()\n\n    # sample_z = th.randn(size=(24, 512))\n\n    scores = []\n    num_iters = 100_000\n    pbar = tqdm(range(num_iters), smoothing=0.1)\n    for i in pbar:\n        vae.train()\n\n        real = next(loader).to(device)\n\n        fake, mu, log_var = vae(real)\n\n        bce = F.binary_cross_entropy(fake, real, size_average=False)\n        kld = -0.5 * th.sum(1 + log_var - mu.pow(2) - log_var.exp())\n        vgg_loss = vgg(fake, real)\n        mse_loss = th.sqrt((fake - real).pow(2).mean())\n\n        loss = bce + kld + lambda_vgg * vgg_loss + lambda_mse * mse_loss\n\n        loss_dict = {\n            \"Total\": loss,\n            \"BCE\": bce,\n            \"Kullback Leibler Divergence\": kld,\n            \"MSE\": mse_loss,\n            \"VGG\": vgg_loss,\n        }\n\n        vae.zero_grad()\n        loss.backward()\n        vae_optim.step()\n\n        wandb.log(loss_dict)\n\n        with th.no_grad():\n            if i % int(num_iters / 100) == 0 or i + 1 == num_iters:\n                vae.eval()\n\n                sample, _, _ = vae(sample_imgs.to(device))\n                grid = utils.make_grid(sample, nrow=6, normalize=True, range=(0, 1))\n                del sample\n                wandb.log({\"Reconstructed Images VAE\": [wandb.Image(grid, caption=f\"Step {i}\")]})\n\n                sample = vae.sampling()\n                grid = utils.make_grid(sample, nrow=6, normalize=True, range=(0, 1))\n                del sample\n                wandb.log({\"Generated Images VAE\": [wandb.Image(grid, caption=f\"Step {i}\")]})\n\n                gc.collect()\n                th.cuda.empty_cache()\n\n                th.save(\n                    {\"vae\": vae.state_dict(), \"vae_optim\": vae_optim.state_dict()},\n                    f\"/home/hans/modelzoo/maua-sg2/vae-{name}-{wandb.run.dir.split('/')[-1].split('-')[-1]}.pt\",\n                )\n\n        if th.isnan(loss).any() or th.isinf(loss).any():\n            print(\"NaN losses, exiting...\")\n            print(\n                {\n                    \"Total\": loss,\n                    \"\\nBCE\": bce,\n                    \"\\nKullback Leibler Divergence\": kld,\n                    \"\\nMSE\": mse_loss,\n                    \"\\nVGG\": vgg_loss,\n                }\n            )\n            wandb.log({\"Total\": 27000})\n            return\n\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser()\n    parser.add_argument(\"--latent_dim\", type=float, default=512)\n    parser.add_argument(\"--num_repeats\", type=float, default=1)\n    parser.add_argument(\"--learning_rate\", type=float, default=0.005)\n    parser.add_argument(\"--lambda_vgg\", type=float, default=1.0)\n    parser.add_argument(\"--lambda_mse\", type=float, default=1.0)\n    args = parser.parse_args()\n\n    device = \"cuda\"\n    th.backends.cudnn.benchmark = True\n\n    wandb.init(project=f\"maua-stylegan\")\n\n    train(\n        args.latent_dim, args.num_repeats, args.learning_rate, args.lambda_vgg, args.lambda_mse,\n    )\n\n"
  },
  {
    "path": "accelerate/accelerate_logcosh.py",
    "content": "import os\nimport gc\nimport wandb\nimport argparse\nimport validation\nimport torch as th\nfrom tqdm import tqdm\nfrom torch.utils import data\nfrom autoencoder import LogCoshVAE\nfrom dataset import MultiResolutionDataset\nfrom torchvision import transforms, utils, models\n\n\ndef data_sampler(dataset, shuffle, distributed):\n    if distributed:\n        return data.distributed.DistributedSampler(dataset, shuffle=shuffle)\n    if shuffle:\n        return data.RandomSampler(dataset)\n    else:\n        return data.SequentialSampler(dataset)\n\n\ndef sample_data(loader):\n    while True:\n        for batch in loader:\n            yield batch\n\n\nclass VGG19(th.nn.Module):\n    \"\"\"\n    Adapted from https://github.com/NVIDIA/pix2pixHD\n    See LICENSE-VGG\n    \"\"\"\n\n    def __init__(self, requires_grad=False):\n        super(VGG19, self).__init__()\n        vgg_pretrained_features = models.vgg19(pretrained=True).features\n        self.slice1 = th.nn.Sequential()\n        self.slice2 = th.nn.Sequential()\n        self.slice3 = th.nn.Sequential()\n        self.slice4 = th.nn.Sequential()\n        self.slice5 = th.nn.Sequential()\n        for x in range(2):\n            self.slice1.add_module(str(x), vgg_pretrained_features[x])\n        for x in range(2, 7):\n            self.slice2.add_module(str(x), vgg_pretrained_features[x])\n        for x in range(7, 12):\n            self.slice3.add_module(str(x), vgg_pretrained_features[x])\n        for x in range(12, 21):\n            self.slice4.add_module(str(x), vgg_pretrained_features[x])\n        for x in range(21, 30):\n            self.slice5.add_module(str(x), vgg_pretrained_features[x])\n        if not requires_grad:\n            for param in self.parameters():\n                param.requires_grad = False\n\n    def forward(self, X):\n        h_relu1 = self.slice1(X)\n        h_relu2 = self.slice2(h_relu1)\n        h_relu3 = self.slice3(h_relu2)\n        h_relu4 = self.slice4(h_relu3)\n        h_relu5 = self.slice5(h_relu4)\n        out = [h_relu1, h_relu2, h_relu3, h_relu4, h_relu5]\n        return out\n\n\nclass VGGLoss(th.nn.Module):\n    \"\"\"\n    Adapted from https://github.com/NVIDIA/pix2pixHD\n    See LICENSE-VGG\n    \"\"\"\n\n    def __init__(self):\n        super(VGGLoss, self).__init__()\n        self.vgg = VGG19().cuda()\n        self.criterion = th.nn.L1Loss()\n        self.weights = [1.0 / 32, 1.0 / 16, 1.0 / 8, 1.0 / 4, 1.0]\n\n    def forward(self, x, y):\n        x_vgg, y_vgg = self.vgg(x), self.vgg(y)\n        loss = 0\n        for i in range(len(x_vgg)):\n            loss += self.weights[i] * self.criterion(x_vgg[i], y_vgg[i].detach())\n        return loss\n\n\ndevice = \"cuda\"\nth.backends.cudnn.benchmark = True\n\nwandb.init(project=f\"maua-stylegan\")\n\n\ndef train(latent_dim, learning_rate, number_filters, vae_alpha, vae_beta, kl_divergence_weight):\n    print(\n        f\"latent_dim={latent_dim}\",\n        f\"learning_rate={learning_rate}\",\n        f\"number_filters={number_filters}\",\n        f\"vae_alpha={vae_alpha}\",\n        f\"vae_beta={vae_beta}\",\n        f\"kl_divergence_weight={kl_divergence_weight}\",\n    )\n\n    batch_size = 64\n    i = None\n    while batch_size >= 1:\n        try:\n            transform = transforms.Compose(\n                [\n                    transforms.Resize(128),\n                    transforms.RandomHorizontalFlip(p=0.5),\n                    transforms.ToTensor(),\n                    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True),\n                ]\n            )\n            data_path = \"/home/hans/trainsets/cyphis\"\n            name = os.path.splitext(os.path.basename(data_path))[0]\n            dataset = MultiResolutionDataset(data_path, transform, 256)\n            dataloader = data.DataLoader(\n                dataset,\n                batch_size=int(batch_size),\n                sampler=data_sampler(dataset, shuffle=True, distributed=False),\n                num_workers=12,\n                drop_last=True,\n            )\n            loader = sample_data(dataloader)\n            sample_imgs = next(loader)[:24]\n            wandb.log(\n                {\"Real Images\": [wandb.Image(utils.make_grid(sample_imgs, nrow=6, normalize=True, range=(-1, 1)))]}\n            )\n\n            hidden_dims = [min(int(number_filters) * 2 ** i, latent_dim) for i in range(5)] + [latent_dim]\n            vae, vae_optim = None, None\n            vae = LogCoshVAE(\n                3, latent_dim, hidden_dims=hidden_dims, alpha=vae_alpha, beta=vae_beta, kld_weight=kl_divergence_weight,\n            ).to(device)\n            vae.train()\n            vae_optim = th.optim.Adam(vae.parameters(), lr=learning_rate)\n\n            mse_loss = th.nn.MSELoss()\n            vgg = VGGLoss()\n\n            sample_z = th.randn(size=(24, latent_dim))\n\n            scores = []\n            num_iters = 100_000\n            pbar = range(num_iters)\n            pbar = tqdm(pbar, smoothing=0.1)\n            for i in pbar:\n                vae.train()\n\n                real = next(loader).to(device)\n                fake, mu, log_var = vae(real)\n\n                loss_dict = vae.loss(real, fake, mu, log_var)\n                vgg_loss = vgg(fake, real)\n                loss = loss_dict[\"Total\"] + vgg_loss\n\n                vae.zero_grad()\n                loss.backward()\n                vae_optim.step()\n\n                wandb.log(\n                    {\n                        \"Total\": loss,\n                        \"VGG\": vgg_loss,\n                        \"Reconstruction\": loss_dict[\"Reconstruction\"],\n                        \"Kullback Leibler Divergence\": loss_dict[\"Kullback Leibler Divergence\"],\n                    }\n                )\n\n                if i % int(num_iters / 1000) == 0 or i + 1 == num_iters:\n                    with th.no_grad():\n                        vae.eval()\n\n                        sample, _, _ = vae(sample_imgs.to(device))\n                        grid = utils.make_grid(sample, nrow=6, normalize=True, range=(-1, 1),)\n                        del sample\n                        wandb.log({\"Reconstructed Images VAE\": [wandb.Image(grid, caption=f\"Step {i}\")]})\n\n                        sample = vae.decode(sample_z.to(device))\n                        grid = utils.make_grid(sample, nrow=6, normalize=True, range=(-1, 1),)\n                        del sample\n                        wandb.log({\"Generated Images VAE\": [wandb.Image(grid, caption=f\"Step {i}\")]})\n\n                if i % int(num_iters / 40) == 0 or i + 1 == num_iters:\n                    with th.no_grad():\n                        fid_dict = validation.vae_fid(vae, int(batch_size), (latent_dim,), 5000, name)\n                        wandb.log(fid_dict)\n                        mse = mse_loss(fake, real) * 5000\n                        score = fid_dict[\"FID\"] + mse + 1000 * vgg_loss\n                        wandb.log({\"Score\": score})\n                        pbar.set_description(f\"FID: {fid_dict['FID']:.2f} MSE: {mse:.2f} VGG: {1000 * vgg_loss:.2f}\")\n\n                    if i >= num_iters / 2:\n                        scores.append(score)\n\n                if th.isnan(loss).any() or th.isinf(loss).any():\n                    print(\"NaN losses, exiting...\")\n                    print(\n                        {\n                            \"Total\": loss.detach().cpu().item(),\n                            \"\\nVGG\": vgg_loss.detach().cpu().item(),\n                            \"\\nReconstruction\": loss_dict[\"Reconstruction\"].detach().cpu().item(),\n                            \"\\nKullback Leibler Divergence\": loss_dict[\"Kullback Leibler Divergence\"]\n                            .detach()\n                            .cpu()\n                            .item(),\n                        }\n                    )\n                    wandb.log({\"Score\": 27000})\n                    return\n\n            return\n\n        except RuntimeError as e:\n            if \"CUDA out of memory\" in str(e):\n                batch_size = batch_size / 2\n\n                if batch_size < 1:\n                    print(\"This configuration does not fit into memory, exiting...\")\n                    wandb.log({\"Score\": 27000})\n                    return\n\n                print(f\"Out of memory, halving batch size... {batch_size}\")\n                if vae is not None:\n                    del vae\n                if vae_optim is not None:\n                    del vae_optim\n                gc.collect()\n                th.cuda.empty_cache()\n\n            else:\n                print(e)\n                return\n\n\nparser = argparse.ArgumentParser()\nparser.add_argument(\"--latent_dim\", type=int, default=1024)\nparser.add_argument(\"--learning_rate\", type=float, default=0.005)\nparser.add_argument(\"--number_filters\", type=int, default=64)\nparser.add_argument(\"--vae_alpha\", type=float, default=10.0)\nparser.add_argument(\"--vae_beta\", type=float, default=1.0)\nparser.add_argument(\"--kl_divergence_weight\", type=float, default=1.0)\nargs = parser.parse_args()\n\ntrain(\n    args.latent_dim, args.learning_rate, args.number_filters, args.vae_alpha, args.vae_beta, args.kl_divergence_weight,\n)\n\n"
  },
  {
    "path": "accelerate/accelerate_segnet.py",
    "content": "import os\nimport gc\nimport wandb\nimport argparse\nimport torch as th\nfrom tqdm import tqdm\nfrom torch.utils import data\nfrom autoencoder import ConvSegNet\nfrom dataset import MultiResolutionDataset\nfrom torchvision import transforms, utils, models\n\n\ndef info(x):\n    print(x.shape, x.detach().cpu().min(), x.detach().cpu().mean(), x.detach().cpu().max())\n\n\ndef sample_data(loader):\n    while True:\n        for batch in loader:\n            yield batch\n\n\nclass VGG19(th.nn.Module):\n    \"\"\"\n    Adapted from https://github.com/NVIDIA/pix2pixHD\n    See LICENSE-VGG\n    \"\"\"\n\n    def __init__(self, requires_grad=False):\n        super(VGG19, self).__init__()\n        vgg_pretrained_features = models.vgg19(pretrained=True).features\n        self.slice1 = th.nn.Sequential()\n        self.slice2 = th.nn.Sequential()\n        self.slice3 = th.nn.Sequential()\n        self.slice4 = th.nn.Sequential()\n        self.slice5 = th.nn.Sequential()\n        for x in range(2):\n            self.slice1.add_module(str(x), vgg_pretrained_features[x])\n        for x in range(2, 7):\n            self.slice2.add_module(str(x), vgg_pretrained_features[x])\n        for x in range(7, 12):\n            self.slice3.add_module(str(x), vgg_pretrained_features[x])\n        for x in range(12, 21):\n            self.slice4.add_module(str(x), vgg_pretrained_features[x])\n        for x in range(21, 30):\n            self.slice5.add_module(str(x), vgg_pretrained_features[x])\n        if not requires_grad:\n            for param in self.parameters():\n                param.requires_grad = False\n\n    def forward(self, X):\n        h_relu1 = self.slice1(X)\n        h_relu2 = self.slice2(h_relu1)\n        h_relu3 = self.slice3(h_relu2)\n        h_relu4 = self.slice4(h_relu3)\n        h_relu5 = self.slice5(h_relu4)\n        out = [h_relu1, h_relu2, h_relu3, h_relu4, h_relu5]\n        return out\n\n\nclass VGGLoss(th.nn.Module):\n    \"\"\"\n    Adapted from https://github.com/NVIDIA/pix2pixHD\n    See LICENSE-VGG\n    \"\"\"\n\n    def __init__(self):\n        super(VGGLoss, self).__init__()\n        self.vgg = VGG19().cuda()\n        self.criterion = th.nn.L1Loss()\n        self.weights = [1.0 / 32, 1.0 / 16, 1.0 / 8, 1.0 / 4, 1.0]\n\n    def forward(self, x, y):\n        x_vgg, y_vgg = self.vgg(x), self.vgg(y)\n        loss = 0\n        for i in range(len(x_vgg)):\n            loss += self.weights[i] * self.criterion(x_vgg[i], y_vgg[i].detach())\n        return loss\n\n\ndef align(x, y, alpha=2):\n    return (x - y).norm(p=2, dim=1).pow(alpha).mean()\n\n\ndef uniform(x, t=2):\n    return (th.pdist(x.view(x.size(0), -1), p=2).pow(2).mul(-t).exp().mean() + 1e-27).log()\n\n\ndef train(learning_rate, lambda_mse):\n    print(\n        f\"learning_rate={learning_rate:.4f}\", f\"lambda_mse={lambda_mse:.4f}\",\n    )\n\n    transform = transforms.Compose(\n        [\n            transforms.Resize(128),\n            transforms.RandomHorizontalFlip(p=0.5),\n            transforms.ToTensor(),\n            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True),\n        ]\n    )\n    batch_size = 72\n    data_path = \"/home/hans/trainsets/cyphis\"\n    name = os.path.splitext(os.path.basename(data_path))[0]\n    dataset = MultiResolutionDataset(data_path, transform, 256)\n    dataloader = data.DataLoader(\n        dataset, batch_size=batch_size, sampler=data.RandomSampler(dataset), num_workers=12, drop_last=True,\n    )\n    loader = sample_data(dataloader)\n    sample_imgs = next(loader)[:24]\n    wandb.log({\"Real Images\": [wandb.Image(utils.make_grid(sample_imgs, nrow=6, normalize=True, range=(-1, 1)))]})\n\n    vae, vae_optim = None, None\n    vae = ConvSegNet().to(device)\n    vae_optim = th.optim.Adam(vae.parameters(), lr=learning_rate)\n\n    vgg = VGGLoss()\n\n    sample_z = th.randn(size=(24, 512, 16, 16))\n    sample_z /= sample_z.abs().max()\n\n    scores = []\n    num_iters = 100_000\n    pbar = tqdm(range(num_iters), smoothing=0.1)\n    for i in pbar:\n        vae.train()\n\n        real = next(loader).to(device)\n\n        z = vae.encode(real)\n        fake = vae.decode(z)\n\n        vgg_loss = vgg(fake, real)\n\n        mse_loss = th.sqrt((fake - real).pow(2).mean())\n\n        # diff = fake - real\n        # recons_loss = recons_alpha * diff + th.log(1.0 + th.exp(-2 * recons_alpha * diff)) - th.log(th.tensor(2.0))\n        # recons_loss = (1.0 / recons_alpha) * recons_loss.mean()\n        # recons_loss = recons_loss if not th.isinf(recons_loss).any() else 0\n\n        # x, y = z.chunk(2)\n        # align_loss = align(x, y, alpha=align_alpha)\n        # unif_loss = -(uniform(x, t=unif_t) + uniform(y, t=unif_t)) / 2.0\n\n        loss = (\n            vgg_loss\n            + lambda_mse * mse_loss\n            # + lambda_recons * recons_loss\n            # + lambda_align * align_loss\n            # + lambda_unif * unif_loss\n        )\n        # print(vgg_loss.detach().cpu().item())\n        # print(lambda_mse * mse_loss.detach().cpu().item())\n        # # print(lambda_recons * recons_loss.detach().cpu().item())\n        # print(lambda_align * align_loss.detach().cpu().item())\n        # print(lambda_unif * unif_loss.detach().cpu().item())\n\n        loss_dict = {\n            \"Total\": loss,\n            \"MSE\": mse_loss,\n            \"VGG\": vgg_loss,\n            # \"Reconstruction\": recons_loss,\n            # \"Alignment\": align_loss,\n            # \"Uniformity\": unif_loss,\n        }\n\n        vae.zero_grad()\n        loss.backward()\n        vae_optim.step()\n\n        wandb.log(loss_dict)\n        # pbar.set_description(\" \".join())\n\n        with th.no_grad():\n            if i % int(num_iters / 100) == 0 or i + 1 == num_iters:\n                vae.eval()\n\n                sample = vae(sample_imgs.to(device))\n                grid = utils.make_grid(sample, nrow=6, normalize=True, range=(-1, 1))\n                del sample\n                wandb.log({\"Reconstructed Images VAE\": [wandb.Image(grid, caption=f\"Step {i}\")]})\n\n                sample = vae.decode(sample_z.to(device))\n                grid = utils.make_grid(sample, nrow=6, normalize=True, range=(-1, 1))\n                del sample\n                wandb.log({\"Generated Images VAE\": [wandb.Image(grid, caption=f\"Step {i}\")]})\n\n                gc.collect()\n                th.cuda.empty_cache()\n\n                th.save(\n                    {\"vae\": vae.state_dict(), \"vae_optim\": vae_optim.state_dict()},\n                    f\"/home/hans/modelzoo/maua-sg2/vae-{name}-{wandb.run.dir.split('/')[-1].split('-')[-1]}.pt\",\n                )\n\n        if th.isnan(loss).any():\n            print(\"NaN losses, exiting...\")\n            wandb.log({\"Total\": 27000})\n            return\n\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser()\n    parser.add_argument(\"--learning_rate\", type=float, default=0.005)\n    parser.add_argument(\"--lambda_mse\", type=float, default=1.0)\n    # parser.add_argument(\"--lambda_recons\", type=float, default=0.0)\n    # parser.add_argument(\"--recons_alpha\", type=float, default=5.0)\n    # parser.add_argument(\"--lambda_align\", type=float, default=1.0)\n    # parser.add_argument(\"--align_alpha\", type=float, default=2.0)\n    # parser.add_argument(\"--lambda_unif\", type=float, default=1.0)\n    # parser.add_argument(\"--unif_t\", type=float, default=0.001)\n    args = parser.parse_args()\n\n    device = \"cuda\"\n    th.backends.cudnn.benchmark = True\n\n    wandb.init(project=f\"maua-stylegan\")\n\n    train(\n        args.learning_rate,\n        args.lambda_mse,\n        # args.lambda_recons,\n        # args.recons_alpha,\n        # args.lambda_align,\n        # args.align_alpha,\n        # args.lambda_unif,\n        # args.unif_t,\n    )\n\n"
  },
  {
    "path": "audioreactive/__init__.py",
    "content": "from .bend import *\nfrom .examples import *\nfrom .latent import *\nfrom .signal import *\nfrom .util import *\n"
  },
  {
    "path": "audioreactive/bend.py",
    "content": "import math\n\nimport kornia.augmentation as kA\nimport kornia.geometry.transform as kT\nimport torch as th\n\n# ====================================================================================\n# ================================= network bending ==================================\n# ====================================================================================\n\n\nclass NetworkBend(th.nn.Module):\n    \"\"\"Base network bending class\n\n    Args:\n        sequential_fn (function): Function that takes a batch of modulation and creates th.nn.Sequential\n        modulation (th.tensor): Modulation batch\n    \"\"\"\n\n    def __init__(self, sequential_fn, modulation):\n        super(NetworkBend, self).__init__()\n        self.sequential = sequential_fn(modulation)\n\n    def forward(self, x):\n        return self.sequential(x)\n\n\nclass AddNoise(th.nn.Module):\n    \"\"\"Adds static noise to output\n\n    Args:\n        noise (th.tensor): Noise to be added\n    \"\"\"\n\n    def __init__(self, noise):\n        super(AddNoise, self).__init__()\n        self.noise = noise\n\n    def forward(self, x):\n        return x + self.noise.to(x.device)\n\n\nclass Print(th.nn.Module):\n    \"\"\"Prints intermediate feature statistics (useful for debugging complicated network bends).\"\"\"\n\n    def forward(self, x):\n        print(x.shape, [x.min().item(), x.mean().item(), x.max().item()], th.std(x).item())\n        return x\n\n\nclass Translate(NetworkBend):\n    \"\"\"Creates horizontal translating effect where repeated linear interpolations from 0 to 1 (saw tooth wave) creates seamless scrolling effect.\n\n    Args:\n        modulation (th.tensor): [0.0-1.0]. Batch of modulation\n        h (int): Height of intermediate features that the network bend is applied to\n        w (int): Width of intermediate features that the network bend is applied to\n        noise (int): Noise to be added (must be 5 * width wide)\n    \"\"\"\n\n    def __init__(self, modulation, h, w, noise):\n        sequential_fn = lambda b: th.nn.Sequential(\n            th.nn.ReflectionPad2d((int(w / 2), int(w / 2), 0, 0)),\n            th.nn.ReflectionPad2d((w, w, 0, 0)),\n            th.nn.ReflectionPad2d((w, 0, 0, 0)),\n            AddNoise(noise),\n            kT.Translate(b),\n            kA.CenterCrop((h, w)),\n        )\n        super(Translate, self).__init__(sequential_fn, modulation)\n\n\nclass Zoom(NetworkBend):\n    \"\"\"Creates zooming effect.\n\n    Args:\n        modulation (th.tensor): [0.0-1.0]. Batch of modulation\n        h (int): height of intermediate features that the network bend is applied to\n        w (int): width of intermediate features that the network bend is applied to\n    \"\"\"\n\n    def __init__(self, modulation, h, w):\n        padding = int(max(h, w)) - 1\n        sequential_fn = lambda b: th.nn.Sequential(th.nn.ReflectionPad2d(padding), kT.Scale(b), kA.CenterCrop((h, w)))\n        super(Zoom, self).__init__(sequential_fn, modulation)\n\n\nclass Rotate(NetworkBend):\n    \"\"\"Creates rotation effect.\n\n    Args:\n        modulation (th.tensor): [0.0-1.0]. Batch of modulation\n        h (int): height of intermediate features that the network bend is applied to\n        w (int): width of intermediate features that the network bend is applied to\n    \"\"\"\n\n    def __init__(self, modulation, h, w):\n        # worst case rotation brings sqrt(2) * max_side_length out-of-frame pixels into frame\n        # padding should cover that exactly\n        padding = int(max(h, w) * (1 - math.sqrt(2) / 2))\n        sequential_fn = lambda b: th.nn.Sequential(th.nn.ReflectionPad2d(padding), kT.Rotate(b), kA.CenterCrop((h, w)))\n        super(Rotate, self).__init__(sequential_fn, modulation)\n"
  },
  {
    "path": "audioreactive/examples/__init__.py",
    "content": "from . import *\n"
  },
  {
    "path": "audioreactive/examples/default.py",
    "content": "import torch as th\n\nimport audioreactive as ar\n\n\ndef initialize(args):\n    args.lo_onsets = ar.onsets(args.audio, args.sr, args.n_frames, fmax=150, smooth=5, clip=97, power=2)\n    args.hi_onsets = ar.onsets(args.audio, args.sr, args.n_frames, fmin=500, smooth=5, clip=99, power=2)\n    return args\n\n\ndef get_latents(selection, args):\n    chroma = ar.chroma(args.audio, args.sr, args.n_frames)\n    chroma_latents = ar.chroma_weight_latents(chroma, selection)\n    latents = ar.gaussian_filter(chroma_latents, 4)\n\n    lo_onsets = args.lo_onsets[:, None, None]\n    hi_onsets = args.hi_onsets[:, None, None]\n\n    latents = hi_onsets * selection[[-4]] + (1 - hi_onsets) * latents\n    latents = lo_onsets * selection[[-7]] + (1 - lo_onsets) * latents\n\n    latents = ar.gaussian_filter(latents, 2, causal=0.2)\n\n    return latents\n\n\ndef get_noise(height, width, scale, num_scales, args):\n    if width > 256:\n        return None\n\n    lo_onsets = args.lo_onsets[:, None, None, None].cuda()\n    hi_onsets = args.hi_onsets[:, None, None, None].cuda()\n\n    noise_noisy = ar.gaussian_filter(th.randn((args.n_frames, 1, height, width), device=\"cuda\"), 5)\n    noise = ar.gaussian_filter(th.randn((args.n_frames, 1, height, width), device=\"cuda\"), 128)\n\n    if width < 128:\n        noise = lo_onsets * noise_noisy + (1 - lo_onsets) * noise\n    if width > 32:\n        noise = hi_onsets * noise_noisy + (1 - hi_onsets) * noise\n\n    noise /= noise.std() * 2.5\n\n    return noise.cpu()\n"
  },
  {
    "path": "audioreactive/examples/kelp.py",
    "content": "\"\"\"\nThis file shows an example of a loop based interpolation\nHere sections are identified with laplacian segmentation and looping latents are generated for each section\nThe noise is looping perlin noise\nLong term section analysis is done with the RMS to interpolate between latent sequences for the intro/outro and drop\n\"\"\"\n\n\nimport librosa as rosa\nimport torch as th\n\nimport audioreactive as ar\n\nOVERRIDE = dict(audio_file=\"audioreactive/examples/Wavefunk - Dwelling in the Kelp.mp3\", out_size=1920)\nBPM = 130\n\n\ndef initialize(args):\n    # RMS can be used to distinguish between the drop sections and intro/outros\n    rms = ar.rms(args.audio, args.sr, args.n_frames, smooth=10, clip=60, power=1)\n    rms = ar.expand(rms, threshold=0.8, ratio=10)\n    rms = ar.gaussian_filter(rms, 4)\n    rms = ar.normalize(rms)\n    args.rms = rms\n\n    # cheating a little here, this my song so I have the multitracks\n    # this is much easier than fiddling with onsets until you have envelopes that dance nicely to the drums\n    audio, sr = rosa.load(\"workspace/kelpkick.wav\", offset=args.offset, duration=args.duration)\n    args.kick_onsets = ar.onsets(audio, sr, args.n_frames, margin=1, smooth=4)\n    audio, sr = rosa.load(\"workspace/kelpsnare.wav\", offset=args.offset, duration=args.duration)\n    args.snare_onsets = ar.onsets(audio, sr, args.n_frames, margin=1, smooth=4)\n\n    ar.plot_signals([args.rms, args.kick_onsets, args.snare_onsets])\n\n    return args\n\n\ndef get_latents(selection, args):\n    # expand envelopes to latent shape\n    rms = args.rms[:, None, None]\n    low_onsets = args.kick_onsets[:, None, None]\n    high_onsets = args.snare_onsets[:, None, None]\n\n    # get timestamps and labels with laplacian segmentation\n    # k is the number of labels the algorithm may use\n    # try multiple values with plot=True to see which value correlates best with the sections of the song\n    timestamps, labels = ar.laplacian_segmentation(args.audio, args.sr, k=7)\n\n    # a second set of latents for the drop section, the 'selection' variable is the other set for the intro\n    drop_selection = ar.load_latents(\"workspace/cyphept_kelp_drop_latents.npy\")\n    color_layer = 9\n\n    latents = []\n    for (start, stop), l in zip(zip(timestamps, timestamps[1:]), labels):\n        start_frame = int(round(start / args.duration * args.n_frames))\n        stop_frame = int(round(stop / args.duration * args.n_frames))\n        section_frames = stop_frame - start_frame\n        section_bars = (stop - start) * (BPM / 60) / 4\n\n        # get portion of latent selection (wrapping around to start)\n        latent_selection_slice = ar.wrapping_slice(selection, l, 4)\n        # spline interpolation loops through selection slice\n        latent_section = ar.spline_loops(latent_selection_slice, n_frames=section_frames, n_loops=section_bars / 4)\n        # set the color with laplacian segmentation label, (1 latent repeated for entire section in upper layers)\n        latent_section[:, color_layer:] = th.cat([selection[[l], color_layer:]] * section_frames)\n\n        # same as above but for the drop latents (with faster loops)\n        drop_selection_slice = ar.wrapping_slice(drop_selection, l, 4)\n        drop_section = ar.spline_loops(drop_selection_slice, n_frames=section_frames, n_loops=section_bars / 2)\n        drop_section[:, color_layer:] = th.cat([drop_selection[[l], color_layer:]] * section_frames)\n\n        # merged based on RMS (drop section or not)\n        latents.append((1 - rms[start_frame:stop_frame]) * latent_section + rms[start_frame:stop_frame] * drop_section)\n\n    # concatenate latents to correct length & smooth over the junctions\n    len_latents = sum([len(l) for l in latents])\n    if len_latents != args.n_frames:\n        latents.append(th.cat([latents[-1][[-1]]] * (args.n_frames - len_latents)))\n    latents = th.cat(latents).float()\n    latents = ar.gaussian_filter(latents, 3)\n\n    # use onsets to modulate towards latents\n    latents = 0.666 * low_onsets * selection[[2]] + (1 - 0.666 * low_onsets) * latents\n    latents = 0.666 * high_onsets * selection[[1]] + (1 - 0.666 * high_onsets) * latents\n\n    latents = ar.gaussian_filter(latents, 1, causal=0.2)\n    return latents\n\n\ndef get_noise(height, width, scale, num_scales, args):\n    if width > 512:  # larger sizes don't fit in VRAM, just use default or randomize\n        return\n\n    num_bars = int(round(args.duration * (BPM / 60) / 4))\n    frames_per_loop = int(args.n_frames / num_bars * 2)  # loop every 2 bars\n\n    def perlin_pls(resolution):\n        perlin = ar.perlin_noise(shape=(frames_per_loop, height, width), res=resolution)[:, None, ...].cpu()\n        perlin = th.cat([perlin] * int(num_bars / 2))  # concatenate multiple copies for looping\n        if args.n_frames - len(perlin) > 0:\n            perlin = th.cat([perlin, th.cat([perlin[[-1]]] * (args.n_frames - len(perlin)))])  # fix up rounding errors\n        return perlin\n\n    smooth = perlin_pls(resolution=(1, 1, 1))  # (time res, x res, y res)\n    noise = perlin_pls(resolution=(8, 4, 4))  # higher resolution => higher frequency noise => more movement in video\n\n    rms = args.rms[:, None, None, None]\n    noise = rms * noise + (1 - rms) * smooth  # blend between noises based on drop (high rms) or not\n\n    return noise\n\n\ndef get_bends(args):\n    # repeat the intermediate features outwards on both sides (2:1 aspect ratio)\n    # + add some noise to give the whole thing a little variation (disguises the repetition)\n    transform = th.nn.Sequential(\n        th.nn.ReplicationPad2d((2, 2, 0, 0)), ar.AddNoise(0.025 * th.randn(size=(1, 1, 4, 8), device=\"cuda\")),\n    )\n    bends = [{\"layer\": 0, \"transform\": transform}]\n\n    return bends\n"
  },
  {
    "path": "audioreactive/examples/tauceti.py",
    "content": "\"\"\"\nThis file shows an example of network bending\nThe latents and noise are similar to temper.py (although without spatial noise controls)\nThe latents cycle through different colors for different sections of the drop\nDuring the drop, a translation is applied which makes the video seem to scroll endlessly\n\"\"\"\n\nfrom functools import partial\n\nimport numpy as np\nimport torch as th\n\nimport audioreactive as ar\n\nOVERRIDE = dict(\n    audio_file=\"audioreactive/examples/Wavefunk - Tau Ceti Alpha.mp3\",\n    out_size=1920,  # get bends assumes 1920x1080 output size\n    dataparallel=False,  # makes use of a kornia transform during network bending => not compatible with dataparallel\n    fps=30,  # 5591 magic number below is based on number of frames in output video with fps of 30\n)\n\n\ndef initialize(args):\n    args.low_onsets = ar.onsets(args.audio, args.sr, args.n_frames, fmax=150, smooth=5, clip=97, power=2)\n    args.high_onsets = ar.onsets(args.audio, args.sr, args.n_frames, fmin=500, smooth=5, clip=99, power=2)\n    return args\n\n\ndef get_latents(selection, args):\n    chroma = ar.chroma(args.audio, args.sr, args.n_frames)\n    chroma_latents = ar.chroma_weight_latents(chroma, selection[:12])  # shape [n_frames, 18, 512]\n    latents = ar.gaussian_filter(chroma_latents, 5)\n\n    lo_onsets = args.low_onsets[:, None, None]  # expand to same shape as latents [n_frames, 1, 1]\n    hi_onsets = args.high_onsets[:, None, None]\n\n    latents = hi_onsets * selection[[-4]] + (1 - hi_onsets) * latents\n    latents = lo_onsets * selection[[-7]] + (1 - lo_onsets) * latents\n\n    latents = ar.gaussian_filter(latents, 5, causal=0)\n\n    # cheating a little, you could probably do this with laplacian segmentation, but is it worth the effort?\n    drop_start = int(5591 * (45 / args.duration))\n    drop_end = int(5591 * (135 / args.duration))\n\n    # selection of latents with different colors (chosen with select_latents.py)\n    color_latent_selection = th.from_numpy(np.load(\"workspace/cyphept-multicolor-latents.npy\"))\n\n    # build sequence of latents for just the upper layers\n    color_layer = 9\n    color_latents = [latents[:drop_start, color_layer:]]\n\n    # for 4 different sections in the drop, use a different color latent\n    drop_length = drop_end - drop_start\n    section_length = int(drop_length / 4)\n    for i, section_start in enumerate(range(0, drop_length, section_length)):\n        if i > 3:\n            break\n        color_latents.append(th.cat([color_latent_selection[[i], color_layer:]] * section_length))\n\n    # ensure color sequence is correct length and concatenate\n    if drop_length - 4 * section_length != 0:\n        color_latents.append(th.cat([color_latent_selection[[i], color_layer:]] * (drop_length - 4 * section_length)))\n    color_latents.append(latents[drop_end:, color_layer:])\n    color_latents = th.cat(color_latents, axis=0)\n\n    color_latents = ar.gaussian_filter(color_latents, 5)\n\n    # set upper layers of latent sequence to the colored sequence\n    latents[:, color_layer:] = color_latents\n\n    return latents\n\n\ndef get_noise(height, width, scale, num_scales, args):\n    if width > 256:\n        return None\n\n    lo_onsets = 1.25 * args.low_onsets[:, None, None, None].cuda()\n    hi_onsets = 1.25 * args.high_onsets[:, None, None, None].cuda()\n\n    noise_noisy = ar.gaussian_filter(th.randn((args.n_frames, 1, height, width), device=\"cuda\"), 5)\n\n    noise = ar.gaussian_filter(th.randn((args.n_frames, 1, height, width), device=\"cuda\"), 128)\n    if width > 8:\n        noise = lo_onsets * noise_noisy + (1 - lo_onsets) * noise\n        noise = hi_onsets * noise_noisy + (1 - hi_onsets) * noise\n\n    noise /= noise.std() * 2.5\n\n    return noise.cpu()\n\n\ndef get_bends(args):\n    # repeat the intermediate features outwards on both sides (2:1 aspect ratio)\n    # + add some noise to give the whole thing a little variation (disguises the repetition)\n    transform = th.nn.Sequential(\n        th.nn.ReplicationPad2d((2, 2, 0, 0)), ar.AddNoise(0.025 * th.randn(size=(1, 1, 4, 8), device=\"cuda\")),\n    )\n    bends = [{\"layer\": 0, \"transform\": transform}]\n\n    # during the drop, create scrolling effect\n    drop_start = int(5591 * (45 / args.duration))\n    drop_end = int(5591 * (135 / args.duration))\n\n    # calculate length of loops, number of loops, and remainder at end of drop\n    scroll_loop_length = int(6 * args.fps)\n    scroll_loop_num = int((drop_end - drop_start) / scroll_loop_length)\n    scroll_trunc = (drop_end - drop_start) - scroll_loop_num * scroll_loop_length\n\n    # apply network bending to 4th layer in StyleGAN\n    # lower layer network bends have more fluid outcomes\n    tl = 4\n    h = 2 ** tl\n    w = 2 * h\n\n    # create values between 0 and 1 corresponding to fraction of scroll from left to right completed\n    # all 0s during intro\n    intro_tl8 = np.zeros(drop_start)\n    # repeating linear interpolation from 0 to 1 during drop\n    loops_tl8 = np.concatenate([np.linspace(0, w, scroll_loop_length)] * scroll_loop_num)\n    # truncated interp\n    last_loop_tl8 = np.linspace(0, w, scroll_loop_length)[:scroll_trunc]\n    # static at final truncated value during outro\n    outro_tl8 = np.ones(args.n_frames - drop_end) * np.linspace(0, w, scroll_loop_length)[scroll_trunc + 1]\n\n    # create 2D array of translations in x and y directions\n    x_tl8 = np.concatenate([intro_tl8, loops_tl8, last_loop_tl8, outro_tl8])\n    y_tl8 = np.zeros(args.n_frames)\n    translation = (th.tensor([x_tl8, y_tl8]).float().T)[: args.n_frames]\n\n    # smooth the transition from intro to drop to prevent jerk\n    translation.T[0, drop_start - args.fps : drop_start + args.fps] = ar.gaussian_filter(\n        translation.T[0, drop_start - 5 * args.fps : drop_start + 5 * args.fps], 5\n    )[4 * args.fps : -4 * args.fps]\n\n    class Translate(NetworkBend):\n        \"\"\"From audioreactive/examples/bend.py\"\"\"\n\n        def __init__(self, modulation, h, w, noise):\n            sequential_fn = lambda b: th.nn.Sequential(\n                th.nn.ReflectionPad2d((int(w / 2), int(w / 2), 0, 0)),  #  < Reflect out to 5x width (so that after\n                th.nn.ReflectionPad2d((w, w, 0, 0)),  #                    < translating w pixels, center crop gives\n                th.nn.ReflectionPad2d((w, 0, 0, 0)),  #                    < same features as translating 0 pixels)\n                AddNoise(noise),  # add some noise to disguise reflections\n                kT.Translate(b),\n                kA.CenterCrop((h, w)),\n            )\n            super(Translate, self).__init__(sequential_fn, modulation)\n\n    # create static noise for translate bend\n    noise = 0.2 * th.randn((1, 1, h, 5 * w), device=\"cuda\")\n    # create function which returns an initialized Translate object when fed a batch of modulation\n    # this is so that creation of the object is delayed until the specific batch is sent into the generator\n    # (there's probably an easier way to do this without the kornia transforms, e.g. using Broad et al.'s transform implementations)\n    transform = lambda batch: partial(Translate, h=h, w=w, noise=noise)(batch)\n    bends += [{\"layer\": tl, \"transform\": transform, \"modulation\": translation}]  # add network bend to list dict\n\n    return bends\n"
  },
  {
    "path": "audioreactive/examples/temper.py",
    "content": "\"\"\"\nThis file shows an example of spatial control of the noise using a simple circular mask\nThe latents are a chromagram weighted sequence, modulated by drum onsets\n\"\"\"\n\nimport scipy.ndimage.filters as ndi\nimport torch as th\n\nimport audioreactive as ar\n\nOVERRIDE = dict(audio_file=\"audioreactive/examples/Wavefunk - Temper.mp3\", out_size=1024)\n\n\ndef initialize(args):\n    # these onsets can definitely use some tweaking, the drum reactivity isn't great for this one\n    # the main bass makes it hard to identify both the kick and the snare because it is so loud and covers the whole spectrum\n    args.lo_onsets = ar.onsets(args.audio, args.sr, args.n_frames, fmax=150, smooth=5, clip=97, power=2)\n    args.hi_onsets = ar.onsets(args.audio, args.sr, args.n_frames, fmin=500, smooth=5, clip=99, power=2)\n    return args\n\n\ndef get_latents(selection, args):\n    # create chromagram weighted sequence\n    chroma = ar.chroma(args.audio, args.sr, args.n_frames)\n    chroma_latents = ar.chroma_weight_latents(chroma, selection)\n    latents = ar.gaussian_filter(chroma_latents, 4)\n\n    # expand onsets to latent shape\n    lo_onsets = args.lo_onsets[:, None, None]\n    hi_onsets = args.hi_onsets[:, None, None]\n\n    # modulate latents to specific latent vectors\n    latents = hi_onsets * selection[[-4]] + (1 - hi_onsets) * latents\n    latents = lo_onsets * selection[[-7]] + (1 - lo_onsets) * latents\n\n    latents = ar.gaussian_filter(latents, 2, causal=0.2)\n\n    return latents\n\n\ndef circular_mask(h, w, center=None, radius=None, soft=0):\n    if center is None:  # use the middle of the image\n        center = (int(w / 2), int(h / 2))\n    if radius is None:  # use the smallest distance between the center and image walls\n        radius = min(center[0], center[1], w - center[0], h - center[1])\n\n    Y, X = np.ogrid[:h, :w]\n    dist_from_center = np.sqrt((X - center[0]) ** 2 + (Y - center[1]) ** 2)\n    mask = dist_from_center <= radius\n\n    if soft > 0:\n        mask = ndi.gaussian_filter(mask, sigma=int(round(soft)))  # blur mask for smoother transition\n\n    return th.from_numpy(mask)\n\n\ndef get_noise(height, width, scale, num_scales, args):\n    if width > 256:  # larger sizes don't fit in VRAM, just use default or randomize\n        return None\n\n    # expand onsets to noise shape\n    # send to GPU as gaussian_filter on large noise tensors with high standard deviation is slow\n    lo_onsets = args.lo_onsets[:, None, None, None].cuda()\n    hi_onsets = args.hi_onsets[:, None, None, None].cuda()\n\n    # 1s inside circle of radius, 0s outside\n    mask = circular_mask(height, width, radius=int(width / 2), soft=2)[None, None, ...].float().cuda()\n\n    # create noise which changes quickly (small standard deviation smoothing)\n    noise_noisy = ar.gaussian_filter(th.randn((args.n_frames, 1, height, width), device=\"cuda\"), 5)\n\n    # create noise which changes slowly (large standard deviation smoothing)\n    noise = ar.gaussian_filter(th.randn((args.n_frames, 1, height, width), device=\"cuda\"), 128)\n\n    # for lower layers, noise inside circle are affected by low onsets\n    if width < 128:\n        noise = 2 * mask * lo_onsets * noise_noisy + (1 - mask) * (1 - lo_onsets) * noise\n    # for upper layers, noise outside circle are affected by high onsets\n    if width > 32:\n        noise = 0.75 * (1 - mask) * hi_onsets * noise_noisy + mask * (1 - 0.75 * hi_onsets) * noise\n\n    # ensure amplitude of noise is close to standard normal distribution (dividing by std. dev. gets it exactly there)\n    noise /= noise.std() * 2\n\n    return noise.cpu()\n"
  },
  {
    "path": "audioreactive/latent.py",
    "content": "import gc\n\nimport numpy as np\nimport torch as th\nfrom scipy import interpolate\n\nfrom models.stylegan2 import Generator\nfrom .signal import gaussian_filter\n\n# ====================================================================================\n# ================================= latent/noise ops =================================\n# ====================================================================================\n\n\ndef chroma_weight_latents(chroma, latents):\n    \"\"\"Creates chromagram weighted latent sequence\n\n    Args:\n        chroma (th.tensor): Chromagram\n        latents (th.tensor): Latents (must have same number as number of notes in chromagram)\n\n    Returns:\n        th.tensor: Chromagram weighted latent sequence\n    \"\"\"\n    base_latents = (chroma[..., None, None] * latents[None, ...]).sum(1)\n    return base_latents\n\n\ndef slerp(val, low, high):\n    \"\"\"Interpolation along geodesic of n-dimensional unit sphere\n    from https://github.com/soumith/dcgan.torch/issues/14#issuecomment-200025792\n\n    Args:\n        val (float): Value between 0 and 1 representing fraction of interpolation completed\n        low (float): Starting value\n        high (float): Ending value\n\n    Returns:\n        float: Interpolated value\n    \"\"\"\n    omega = np.arccos(np.clip(np.dot(low / np.linalg.norm(low), high / np.linalg.norm(high)), -1, 1))\n    so = np.sin(omega)\n    if so == 0:\n        return (1.0 - val) * low + val * high  # L'Hopital's rule/LERP\n    return np.sin((1.0 - val) * omega) / so * low + np.sin(val * omega) / so * high\n\n\ndef slerp_loops(latent_selection, n_frames, n_loops, smoothing=1, loop=True):\n    \"\"\"Get looping latents using geodesic interpolation. Total length of n_frames with n_loops repeats.\n\n    Args:\n        latent_selection (th.tensor): Set of latents to loop between (in order)\n        n_frames (int): Total length of output looping sequence\n        n_loops (int): Number of times to loop\n        smoothing (int, optional): Standard deviation of gaussian smoothing kernel. Defaults to 1.\n        loop (bool, optional): Whether to return to first latent. Defaults to True.\n\n    Returns:\n        th.tensor: Sequence of smoothly looping latents\n    \"\"\"\n    if loop:\n        latent_selection = np.concatenate([latent_selection, latent_selection[[0]]])\n\n    base_latents = []\n    for n in range(len(latent_selection)):\n        for val in np.linspace(0.0, 1.0, int(n_frames // max(1, n_loops) // len(latent_selection))):\n            base_latents.append(\n                th.from_numpy(\n                    slerp(\n                        val,\n                        latent_selection[n % len(latent_selection)][0],\n                        latent_selection[(n + 1) % len(latent_selection)][0],\n                    )\n                )\n            )\n    base_latents = th.stack(base_latents)\n    base_latents = gaussian_filter(base_latents, smoothing)\n    base_latents = th.cat([base_latents] * int(n_frames / len(base_latents)), axis=0)\n    base_latents = th.cat([base_latents[:, None, :]] * 18, axis=1)\n    if n_frames - len(base_latents) != 0:\n        base_latents = th.cat([base_latents, base_latents[0 : n_frames - len(base_latents)]])\n    return base_latents\n\n\ndef spline_loops(latent_selection, n_frames, n_loops, loop=True):\n    \"\"\"Get looping latents using spline interpolation. Total length of n_frames with n_loops repeats.\n\n    Args:\n        latent_selection (th.tensor): Set of latents to loop between (in order)\n        n_frames (int): Total length of output looping sequence\n        n_loops (int): Number of times to loop\n        loop (bool, optional): Whether to return to first latent. Defaults to True.\n\n    Returns:\n        th.tensor: Sequence of smoothly looping latents\n    \"\"\"\n    if loop:\n        latent_selection = np.concatenate([latent_selection, latent_selection[[0]]])\n\n    x = np.linspace(0, 1, int(n_frames // max(1, n_loops)))\n    base_latents = np.zeros((len(x), *latent_selection.shape[1:]))\n    for lay in range(latent_selection.shape[1]):\n        for lat in range(latent_selection.shape[2]):\n            tck = interpolate.splrep(np.linspace(0, 1, latent_selection.shape[0]), latent_selection[:, lay, lat])\n            base_latents[:, lay, lat] = interpolate.splev(x, tck)\n\n    base_latents = th.cat([th.from_numpy(base_latents)] * int(n_frames / len(base_latents)), axis=0)\n    if n_frames - len(base_latents) > 0:\n        base_latents = th.cat([base_latents, base_latents[0 : n_frames - len(base_latents)]])\n    return base_latents[:n_frames]\n\n\ndef wrapping_slice(tensor, start, length, return_indices=False):\n    \"\"\"Gets slice of tensor of a given length that wraps around to beginning\n\n    Args:\n        tensor (th.tensor): Tensor to slice\n        start (int): Starting index\n        length (int): Size of slice\n        return_indices (bool, optional): Whether to return indices rather than values. Defaults to False.\n\n    Returns:\n        th.tensor: Values or indices of slice\n    \"\"\"\n    if start + length <= tensor.shape[0]:\n        indices = th.arange(start, start + length)\n    else:\n        indices = th.cat((th.arange(start, tensor.shape[0]), th.arange(0, (start + length) % tensor.shape[0])))\n    if tensor.shape[0] == 1:\n        indices = th.zeros(1, dtype=th.int64)\n    if return_indices:\n        return indices\n    return tensor[indices]\n\n\ndef generate_latents(n_latents, ckpt, G_res, noconst=False, latent_dim=512, n_mlp=8, channel_multiplier=2):\n    \"\"\"Generates random, mapped latents\n\n    Args:\n        n_latents (int): Number of mapped latents to generate \n        ckpt (str): Generator checkpoint to use\n        G_res (int): Generator's training resolution\n        noconst (bool, optional): Whether the generator was trained without constant starting layer. Defaults to False.\n        latent_dim (int, optional): Size of generator's latent vectors. Defaults to 512.\n        n_mlp (int, optional): Number of layers in the generator's mapping network. Defaults to 8.\n        channel_multiplier (int, optional): Scaling multiplier for generator's channel depth. Defaults to 2.\n\n    Returns:\n        th.tensor: Set of mapped latents\n    \"\"\"\n    generator = Generator(\n        G_res, latent_dim, n_mlp, channel_multiplier=channel_multiplier, constant_input=not noconst, checkpoint=ckpt,\n    ).cuda()\n    zs = th.randn((n_latents, latent_dim), device=\"cuda\")\n    latent_selection = generator(zs, map_latents=True).cpu()\n    del generator, zs\n    gc.collect()\n    th.cuda.empty_cache()\n    return latent_selection\n\n\ndef save_latents(latents, filename):\n    \"\"\"Saves latent vectors to file\n\n    Args:\n        latents (th.tensor): Latent vector(s) to save\n        filename (str): Filename to save to\n    \"\"\"\n    np.save(filename, latents)\n\n\ndef load_latents(filename):\n    \"\"\"Load latents from numpy file\n\n    Args:\n        filename (str): Filename to load from\n\n    Returns:\n        th.tensor: Latent vectors\n    \"\"\"\n    return th.from_numpy(np.load(filename))\n\n\ndef _perlinterpolant(t):\n    return t * t * t * (t * (t * 6 - 15) + 10)\n\n\ndef perlin_noise(shape, res, tileable=(True, False, False), interpolant=_perlinterpolant):\n    \"\"\"Generate a 3D tensor of perlin noise.\n\n    Args:\n        shape: The shape of the generated tensor (tuple of three ints). This must be a multiple of res.\n        res: The number of periods of noise to generate along each axis (tuple of three ints). Note shape must be a multiple of res.\n        tileable: If the noise should be tileable along each axis (tuple of three bools). Defaults to (False, False, False).\n        interpolant: The interpolation function, defaults to t*t*t*(t*(t*6 - 15) + 10).\n\n    Returns:\n        A tensor of shape shape with the generated noise.\n\n    Raises:\n        ValueError: If shape is not a multiple of res.\n    \"\"\"\n    delta = (res[0] / shape[0], res[1] / shape[1], res[2] / shape[2])\n    d = (shape[0] // res[0], shape[1] // res[1], shape[2] // res[2])\n    grid = np.mgrid[0 : res[0] : delta[0], 0 : res[1] : delta[1], 0 : res[2] : delta[2]]\n    grid = grid.transpose(1, 2, 3, 0) % 1\n    grid = th.from_numpy(grid).cuda()\n    # Gradients\n    theta = 2 * np.pi * np.random.rand(res[0] + 1, res[1] + 1, res[2] + 1)\n    phi = 2 * np.pi * np.random.rand(res[0] + 1, res[1] + 1, res[2] + 1)\n    gradients = np.stack((np.sin(phi) * np.cos(theta), np.sin(phi) * np.sin(theta), np.cos(phi)), axis=3)\n    if tileable[0]:\n        gradients[-1, :, :] = gradients[0, :, :]\n    if tileable[1]:\n        gradients[:, -1, :] = gradients[:, 0, :]\n    if tileable[2]:\n        gradients[:, :, -1] = gradients[:, :, 0]\n    gradients = gradients.repeat(d[0], 0).repeat(d[1], 1).repeat(d[2], 2)\n    gradients = th.from_numpy(gradients).cuda()\n    g000 = gradients[: -d[0], : -d[1], : -d[2]]\n    g100 = gradients[d[0] :, : -d[1], : -d[2]]\n    g010 = gradients[: -d[0], d[1] :, : -d[2]]\n    g110 = gradients[d[0] :, d[1] :, : -d[2]]\n    g001 = gradients[: -d[0], : -d[1], d[2] :]\n    g101 = gradients[d[0] :, : -d[1], d[2] :]\n    g011 = gradients[: -d[0], d[1] :, d[2] :]\n    g111 = gradients[d[0] :, d[1] :, d[2] :]\n    # Ramps\n    n000 = th.sum(th.stack((grid[:, :, :, 0], grid[:, :, :, 1], grid[:, :, :, 2]), axis=3) * g000, 3)\n    n100 = th.sum(th.stack((grid[:, :, :, 0] - 1, grid[:, :, :, 1], grid[:, :, :, 2]), axis=3) * g100, 3)\n    n010 = th.sum(th.stack((grid[:, :, :, 0], grid[:, :, :, 1] - 1, grid[:, :, :, 2]), axis=3) * g010, 3)\n    n110 = th.sum(th.stack((grid[:, :, :, 0] - 1, grid[:, :, :, 1] - 1, grid[:, :, :, 2]), axis=3) * g110, 3)\n    n001 = th.sum(th.stack((grid[:, :, :, 0], grid[:, :, :, 1], grid[:, :, :, 2] - 1), axis=3) * g001, 3)\n    n101 = th.sum(th.stack((grid[:, :, :, 0] - 1, grid[:, :, :, 1], grid[:, :, :, 2] - 1), axis=3) * g101, 3)\n    n011 = th.sum(th.stack((grid[:, :, :, 0], grid[:, :, :, 1] - 1, grid[:, :, :, 2] - 1), axis=3) * g011, 3)\n    n111 = th.sum(th.stack((grid[:, :, :, 0] - 1, grid[:, :, :, 1] - 1, grid[:, :, :, 2] - 1), axis=3) * g111, 3)\n    # Interpolation\n    t = interpolant(grid)\n    n00 = n000 * (1 - t[:, :, :, 0]) + t[:, :, :, 0] * n100\n    n10 = n010 * (1 - t[:, :, :, 0]) + t[:, :, :, 0] * n110\n    n01 = n001 * (1 - t[:, :, :, 0]) + t[:, :, :, 0] * n101\n    n11 = n011 * (1 - t[:, :, :, 0]) + t[:, :, :, 0] * n111\n    n0 = (1 - t[:, :, :, 1]) * n00 + t[:, :, :, 1] * n10\n    n1 = (1 - t[:, :, :, 1]) * n01 + t[:, :, :, 1] * n11\n    perlin = (1 - t[:, :, :, 2]) * n0 + t[:, :, :, 2] * n1\n    return perlin * 2 - 1  # stretch from -1 to 1\n"
  },
  {
    "path": "audioreactive/signal.py",
    "content": "import os\nimport warnings\nfrom pathlib import Path\n\nimport joblib\nimport librosa as rosa\nimport librosa.display\nimport madmom as mm\nimport matplotlib.patches as patches\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport scipy\nimport scipy.signal as signal\nimport sklearn.cluster\nimport torch as th\nimport torch.nn.functional as F\n\nSMF = 1  # this is set by generate_audiovisual.py based on rendering fps\n\n\ndef set_SMF(smf):\n    global SMF\n    SMF = smf\n\n\n# ====================================================================================\n# ==================================== signal ops ====================================\n# ====================================================================================\n\n\ndef onsets(audio, sr, n_frames, margin=8, fmin=20, fmax=8000, smooth=1, clip=100, power=1, type=\"mm\"):\n    \"\"\"Creates onset envelope from audio\n\n    Args:\n        audio (np.array): Audio signal\n        sr (int): Sampling rate of the audio\n        n_frames (int): Total number of frames to resample envelope to\n        margin (int, optional): For percussive source separation, higher values create more extreme separations. Defaults to 8.\n        fmin (int, optional): Minimum frequency for onset analysis. Defaults to 20.\n        fmax (int, optional): Maximum frequency for onset analysis. Defaults to 8000.\n        smooth (int, optional): Standard deviation of gaussian kernel to smooth with. Defaults to 1.\n        clip (int, optional): Percentile to clip onset signal to. Defaults to 100.\n        power (int, optional): Exponent to raise onset signal to. Defaults to 1.\n        type (str, optional): [\"rosa\", \"mm\"]. Whether to use librosa or madmom for onset analysis. Madmom is slower but often more accurate. Defaults to \"mm\".\n\n    Returns:\n        th.tensor, shape=(n_frames,): Onset envelope\n    \"\"\"\n    y_perc = rosa.effects.percussive(y=audio, margin=margin)\n    if type == \"rosa\":\n        onset = rosa.onset.onset_strength(y=y_perc, sr=sr, fmin=fmin, fmax=fmax)\n    elif type == \"mm\":\n        sig = mm.audio.signal.Signal(y_perc, num_channels=1, sample_rate=sr)\n        sig_frames = mm.audio.signal.FramedSignal(sig, frame_size=2048, hop_size=441)\n        stft = mm.audio.stft.ShortTimeFourierTransform(sig_frames, circular_shift=True)\n        spec = mm.audio.spectrogram.Spectrogram(stft, circular_shift=True)\n        filt_spec = mm.audio.spectrogram.FilteredSpectrogram(spec, num_bands=24, fmin=fmin, fmax=fmax)\n        onset = np.sum(\n            [\n                mm.features.onsets.spectral_diff(filt_spec),\n                mm.features.onsets.spectral_flux(filt_spec),\n                mm.features.onsets.superflux(filt_spec),\n                mm.features.onsets.complex_flux(filt_spec),\n                mm.features.onsets.modified_kullback_leibler(filt_spec),\n            ],\n            axis=0,\n        )\n    onset = np.clip(signal.resample(onset, n_frames), onset.min(), onset.max())\n    onset = th.from_numpy(onset).float()\n    onset = gaussian_filter(onset, smooth, causal=0)\n    onset = percentile_clip(onset, clip)\n    onset = onset ** power\n    return onset\n\n\ndef rms(y, sr, n_frames, fmin=20, fmax=8000, smooth=180, clip=50, power=6):\n    \"\"\"Creates RMS envelope from audio\n\n    Args:\n        audio (np.array): Audio signal\n        sr (int): Sampling rate of the audio\n        n_frames (int): Total number of frames to resample envelope to\n        fmin (int, optional): Minimum frequency for onset analysis. Defaults to 20.\n        fmax (int, optional): Maximum frequency for onset analysis. Defaults to 8000.\n        smooth (int, optional): Standard deviation of gaussian kernel to smooth with. Defaults to 180.\n        clip (int, optional): Percentile to clip onset signal to. Defaults to 50.\n        power (int, optional): Exponent to raise onset signal to. Defaults to 6.\n\n    Returns:\n        th.tensor, shape=(n_frames,): RMS envelope\n    \"\"\"\n    y_filt = signal.sosfilt(signal.butter(12, [fmin, fmax], \"bp\", fs=sr, output=\"sos\"), y)\n    rms = rosa.feature.rms(S=np.abs(rosa.stft(y=y_filt, hop_length=512)))[0]\n    rms = np.clip(signal.resample(rms, n_frames), rms.min(), rms.max())\n    rms = th.from_numpy(rms).float()\n    rms = gaussian_filter(rms, smooth, causal=0.05)\n    rms = percentile_clip(rms, clip)\n    rms = rms ** power\n    return rms\n\n\ndef raw_chroma(audio, sr, type=\"cens\", nearest_neighbor=True):\n    \"\"\"Creates chromagram\n\n    Args:\n        audio (np.array): Audio signal\n        sr (int): Sampling rate of the audio\n        type (str, optional): [\"cens\", \"cqt\", \"stft\", \"deep\", \"clp\"]. Which strategy to use to calculate the chromagram. Defaults to \"cens\".\n        nearest_neighbor (bool, optional): Whether to post process using nearest neighbor smoothing. Defaults to True.\n\n    Returns:\n        np.array, shape=(12, n_frames): Chromagram\n    \"\"\"\n    if type == \"cens\":\n        ch = rosa.feature.chroma_cens(y=audio, sr=sr)\n    elif type == \"cqt\":\n        ch = rosa.feature.chroma_cqt(y=audio, sr=sr)\n    elif type == \"stft\":\n        ch = rosa.feature.chroma_stft(y=audio, sr=sr)\n    elif type == \"deep\":\n        sig = mm.audio.signal.Signal(audio, num_channels=1, sample_rate=sr)\n        ch = mm.audio.chroma.DeepChromaProcessor().process(sig).T\n    elif type == \"clp\":\n        sig = mm.audio.signal.Signal(audio, num_channels=1, sample_rate=sr)\n        ch = mm.audio.chroma.CLPChromaProcessor().process(sig).T\n    else:\n        print(\"chroma type not recognized, options are: [cens, cqt, deep, clp, or stft]. defaulting to cens...\")\n        ch = rosa.feature.chroma_cens(y=audio, sr=sr)\n\n    if nearest_neighbor:\n        ch = np.minimum(ch, rosa.decompose.nn_filter(ch, aggregate=np.median, metric=\"cosine\"))\n\n    return ch\n\n\ndef chroma(audio, sr, n_frames, margin=16, type=\"cens\", notes=12):\n    \"\"\"Creates chromagram for the harmonic component of the audio\n\n    Args:\n        audio (np.array): Audio signal\n        sr (int): Sampling rate of the audio\n        n_frames (int): Total number of frames to resample envelope to\n        margin (int, optional): For harmonic source separation, higher values create more extreme separations. Defaults to 16.\n        type (str, optional): [\"cens\", \"cqt\", \"stft\", \"deep\", \"clp\"]. Which strategy to use to calculate the chromagram. Defaults to \"cens\".\n        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.\n\n    Returns:\n        th.tensor, shape=(n_frames, 12): Chromagram\n    \"\"\"\n    y_harm = rosa.effects.harmonic(y=audio, margin=margin)\n    chroma = raw_chroma(y_harm, sr, type=type).T\n    chroma = signal.resample(chroma, n_frames)\n    notes_indices = np.argsort(np.median(chroma, axis=0))[:notes]\n    chroma = chroma[:, notes_indices]\n    chroma = th.from_numpy(chroma / chroma.sum(1)[:, None]).float()\n    return chroma\n\n\ndef laplacian_segmentation(signal, sr, k=5, plot=False):\n    \"\"\"Segments the audio with pattern recurrence analysis\n    From https://librosa.org/doc/latest/auto_examples/plot_segmentation.html#sphx-glr-auto-examples-plot-segmentation-py%22\n\n    Args:\n        signal (np.array): Audio signal\n        sr (int): Sampling rate of the audio\n        k (int, optional): Number of labels to use during segmentation. Defaults to 5.\n        plot (bool, optional): Whether to show plot of found segmentation. Defaults to False.\n\n    Returns:\n        tuple(list, list): List of starting timestamps and labels of found segments\n    \"\"\"\n    BINS_PER_OCTAVE = 12 * 3\n    N_OCTAVES = 7\n    C = librosa.amplitude_to_db(\n        np.abs(librosa.cqt(y=signal, sr=sr, bins_per_octave=BINS_PER_OCTAVE, n_bins=N_OCTAVES * BINS_PER_OCTAVE)),\n        ref=np.max,\n    )\n\n    # make CQT beat-synchronous to reduce dimensionality\n    tempo, beats = librosa.beat.beat_track(y=signal, sr=sr, trim=False)\n    Csync = librosa.util.sync(C, beats, aggregate=np.median)\n\n    # build a weighted recurrence matrix using beat-synchronous CQT\n    R = librosa.segment.recurrence_matrix(Csync, width=3, mode=\"affinity\", sym=True)\n    # enhance diagonals with a median filter\n    df = librosa.segment.timelag_filter(scipy.ndimage.median_filter)\n    Rf = df(R, size=(1, 7))\n\n    # build the sequence matrix using mfcc-similarity\n    mfcc = librosa.feature.mfcc(y=signal, sr=sr)\n    Msync = librosa.util.sync(mfcc, beats)\n    path_distance = np.sum(np.diff(Msync, axis=1) ** 2, axis=0)\n    sigma = np.median(path_distance)\n    path_sim = np.exp(-path_distance / sigma)\n    R_path = np.diag(path_sim, k=1) + np.diag(path_sim, k=-1)\n\n    # compute the balanced combination\n    deg_path = np.sum(R_path, axis=1)\n    deg_rec = np.sum(Rf, axis=1)\n    mu = deg_path.dot(deg_path + deg_rec) / np.sum((deg_path + deg_rec) ** 2)\n\n    A = mu * Rf + (1 - mu) * R_path\n    # compute the normalized laplacian and its spectral decomposition\n    L = scipy.sparse.csgraph.laplacian(A, normed=True)\n    evals, evecs = scipy.linalg.eigh(L)\n    # median filter to smooth over small discontinuities\n    evecs = scipy.ndimage.median_filter(evecs, size=(9, 1))\n    # cumulative normalization for symmetric normalized laplacian eigenvectors\n    Cnorm = np.cumsum(evecs ** 2, axis=1) ** 0.5\n\n    X = evecs[:, :k] / Cnorm[:, k - 1 : k]\n\n    # use first k components to cluster beats into segments\n    seg_ids = sklearn.cluster.KMeans(n_clusters=k).fit_predict(X)\n\n    bound_beats = 1 + np.flatnonzero(seg_ids[:-1] != seg_ids[1:])  # locate segment boundaries from the label sequence\n    bound_beats = librosa.util.fix_frames(bound_beats, x_min=0)  # count beat 0 as a boundary\n    bound_segs = list(seg_ids[bound_beats])  # compute the segment label for each boundary\n    bound_frames = beats[bound_beats]  # convert beat indices to frames\n    bound_frames = librosa.util.fix_frames(bound_frames, x_min=None, x_max=C.shape[1] - 1)\n    bound_times = librosa.frames_to_time(bound_frames)\n    if bound_times[0] != 0:\n        bound_times[0] = 0\n\n    if plot:\n        freqs = librosa.cqt_frequencies(\n            n_bins=C.shape[0], fmin=librosa.note_to_hz(\"C1\"), bins_per_octave=BINS_PER_OCTAVE\n        )\n        fig, ax = plt.subplots()\n        colors = plt.get_cmap(\"Paired\", k)\n        librosa.display.specshow(C, y_axis=\"cqt_hz\", sr=sr, bins_per_octave=BINS_PER_OCTAVE, x_axis=\"time\", ax=ax)\n        for interval, label in zip(zip(bound_times, bound_times[1:]), bound_segs):\n            ax.add_patch(\n                patches.Rectangle(\n                    (interval[0], freqs[0]), interval[1] - interval[0], freqs[-1], facecolor=colors(label), alpha=0.50\n                )\n            )\n        plt.show()\n\n    return list(bound_times), list(bound_segs)\n\n\ndef normalize(signal):\n    \"\"\"Normalize signal between 0 and 1\n\n    Args:\n        signal (np.array/th.tensor): Signal to normalize\n\n    Returns:\n        np.array/th.tensor: Normalized signal\n    \"\"\"\n    signal -= signal.min()\n    signal /= signal.max()\n    return signal\n\n\ndef percentile(signal, p):\n    \"\"\"Calculate percentile of signal\n\n    Args:\n        signal (np.array/th.tensor): Signal to normalize\n        p (int): [0-100]. Percentile to find\n\n    Returns:\n        int: Percentile signal value\n    \"\"\"\n    k = 1 + round(0.01 * float(p) * (signal.numel() - 1))\n    return signal.view(-1).kthvalue(k).values.item()\n\n\ndef percentile_clip(signal, p):\n    \"\"\"Normalize signal between 0 and 1, clipping peak values above given percentile\n\n    Args:\n        signal (th.tensor): Signal to normalize\n        p (int): [0-100]. Percentile to clip to\n\n    Returns:\n        th.tensor: Clipped signal\n    \"\"\"\n    locs = th.arange(0, signal.shape[0])\n    peaks = th.ones(signal.shape, dtype=bool)\n    main = signal.take(locs)\n\n    plus = signal.take((locs + 1).clamp(0, signal.shape[0] - 1))\n    minus = signal.take((locs - 1).clamp(0, signal.shape[0] - 1))\n    peaks &= th.gt(main, plus)\n    peaks &= th.gt(main, minus)\n\n    signal = signal.clamp(0, percentile(signal[peaks], p))\n    signal /= signal.max()\n    return signal\n\n\ndef compress(signal, threshold, ratio, invert=False):\n    \"\"\"Expand or compress signal. Values above/below (depending on invert) threshold are multiplied by ratio.\n\n    Args:\n        signal (th.tensor): Signal to normalize\n        threshold (int): Signal value above/below which to change signal\n        ratio (float): Value to multiply signal with\n        invert (bool, optional): Specifies if values above or below threshold are affected. Defaults to False.\n\n    Returns:\n        th.tensor: Compressed/expanded signal\n    \"\"\"\n    if invert:\n        signal[signal < threshold] *= ratio\n    else:\n        signal[signal > threshold] *= ratio\n    return normalize(signal)\n\n\ndef expand(signal, threshold, ratio, invert=False):\n    \"\"\"Alias for compress. Whether compression or expansion occurs is determined by values of threshold and ratio\"\"\"\n    return compress(signal, threshold, ratio, invert)\n\n\ndef gaussian_filter(x, sigma, causal=None):\n    \"\"\"Smooth tensors along time (first) axis with gaussian kernel.\n\n    Args:\n        x (th.tensor): Tensor to be smoothed\n        sigma (float): Standard deviation for gaussian kernel (higher value gives smoother result)\n        causal (float, optional): Factor to multiply right side of gaussian kernel with. Lower value decreases effect of \"future\" values. Defaults to None.\n\n    Returns:\n        th.tensor: Smoothed tensor\n    \"\"\"\n    dim = len(x.shape)\n    n_frames = x.shape[0]\n    while len(x.shape) < 3:\n        x = x[:, None]\n\n    radius = min(int(sigma * 4 * SMF), 3 * len(x))\n    channels = x.shape[1]\n\n    kernel = th.arange(-radius, radius + 1, dtype=th.float32, device=x.device)\n    kernel = th.exp(-0.5 / sigma ** 2 * kernel ** 2)\n    if causal is not None:\n        kernel[radius + 1 :] *= 0 if not isinstance(causal, float) else causal\n    kernel = kernel / kernel.sum()\n    kernel = kernel.view(1, 1, len(kernel)).repeat(channels, 1, 1)\n\n    if dim == 4:\n        t, c, h, w = x.shape\n        x = x.view(t, c, h * w)\n    x = x.transpose(0, 2)\n\n    if radius > n_frames:  # prevent padding errors on short sequences\n        x = F.pad(x, (n_frames, n_frames), mode=\"circular\")\n        print(\n            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}).\"\n        )\n        x = F.pad(x, (radius - n_frames, radius - n_frames), mode=\"constant\")\n    else:\n        x = F.pad(x, (radius, radius), mode=\"circular\")\n\n    x = F.conv1d(x, weight=kernel, groups=channels)\n\n    x = x.transpose(0, 2)\n    if dim == 4:\n        x = x.view(t, c, h, w)\n\n    if len(x.shape) > dim:\n        x = x.squeeze()\n\n    return x\n\n\ndef load_audio(audio_file, offset=0, duration=-1, cache=True):\n    \"\"\"Handles loading of audio files. Automatically caches to .npy files to increase loading speed.\n\n    Args:\n        audio_file (str): Path to audio file to load\n        offset (float, optional): Time (in seconds) to start from. Defaults to 0.\n        duration (float, optional): Length of time to load in. Defaults to -1 (full duration).\n        cache (bool): Whether to cache the raw audio file or not\n\n    Returns:\n        audio   : audio signal\n        sr      : sample rate of audio\n        duration: duration of audio in seconds\n    \"\"\"\n    audio_dur = rosa.get_duration(filename=audio_file)\n    if duration == -1 or audio_dur < duration:\n        duration = audio_dur\n        if offset != 0:\n            duration -= offset\n\n    cache_file = (\n        f\"workspace/{Path(audio_file).stem}\"\n        + (\"\" if duration == -1 else f\"_length{duration}\")\n        + (\"\" if offset == 0 else f\"_start{offset}\")\n        + \".npy\"\n    )\n    if cache and not os.path.exists(cache_file):\n        with warnings.catch_warnings():\n            warnings.filterwarnings(\"ignore\", message=\"PySoundFile failed. Trying audioread instead.\")\n            audio, sr = rosa.load(audio_file, offset=offset, duration=duration)\n        joblib.dump((audio, sr), cache_file)\n    else:\n        audio, sr = joblib.load(cache_file)\n\n    return audio, sr, duration\n"
  },
  {
    "path": "audioreactive/util.py",
    "content": "import librosa as rosa\nimport librosa.display\nimport matplotlib.pyplot as plt\nimport numpy as np\n\n# ====================================================================================\n# ==================================== utilities =====================================\n# ====================================================================================\n\n\ndef info(arr):\n    \"\"\"Shows statistics and shape information of (lists of) np.arrays/th.tensors\n\n    Args:\n        arr (np.array/th.tensor/list): List of or single np.array or th.tensor\n    \"\"\"\n    if isinstance(arr, list):\n        print([(list(a.shape), f\"{a.min():.2f}\", f\"{a.mean():.2f}\", f\"{a.max():.2f}\") for a in arr])\n    else:\n        print(list(arr.shape), f\"{arr.min():.2f}\", f\"{arr.mean():.2f}\", f\"{arr.max():.2f}\")\n\n\ndef plot_signals(signals):\n    \"\"\"Shows plot of (multiple) 1D signals\n\n    Args:\n        signals (np.array/th.tensor): List of signals (1 non-unit dimension)\n    \"\"\"\n    plt.figure(figsize=(16, 4 * len(signals)))\n    for sbplt, y in enumerate(signals):\n        try:\n            signal = signal.cpu().numpy()\n        except:\n            pass\n        plt.subplot(len(signals), 1, sbplt + 1)\n        plt.plot(y.squeeze())\n    plt.tight_layout()\n    plt.show()\n\n\ndef plot_spectra(spectra, chroma=False):\n    \"\"\"Shows plot of (multiple) spectrograms\n\n    Args:\n        spectra (np.array/th.tensor): List of spectrograms\n        chroma (bool, optional): Whether to plot with chromagram y-axis label. Defaults to False.\n    \"\"\"\n    fig, axes = plt.subplots(len(spectra), 1, figsize=(16, 4 * len(spectra)))\n    for ax, spectrum in zip(axes if len(spectra) > 1 else [axes], spectra):\n        try:\n            spectrum = spectrum.cpu().numpy()\n        except:\n            pass\n        if spectrum.shape[1] == 12:\n            spectrum = spectrum.T\n        rosa.display.specshow(spectrum, y_axis=\"chroma\" if chroma else None, x_axis=\"time\", ax=ax)\n    plt.tight_layout()\n    plt.show()\n\n\ndef plot_audio(audio, sr):\n    \"\"\"Shows spectrogram of audio signal\n\n    Args:\n        audio (np.array): Audio signal to be plotted\n        sr (int): Sampling rate of the audio\n    \"\"\"\n    plt.figure(figsize=(16, 9))\n    rosa.display.specshow(\n        rosa.power_to_db(rosa.feature.melspectrogram(y=audio, sr=sr), ref=np.max), y_axis=\"mel\", x_axis=\"time\"\n    )\n    plt.colorbar(format=\"%+2.f dB\")\n    plt.tight_layout()\n    plt.show()\n\n\ndef plot_chroma_comparison(audio, sr):\n    \"\"\"Shows plot comparing different chromagram strategies.\n\n    Args:\n        audio (np.array): Audio signal to be plotted\n        sr (int): Sampling rate of the audio\n    \"\"\"\n    fig, ax = plt.subplots(nrows=2, ncols=3, figsize=(16, 9))\n    for col, types in enumerate([[\"cens\", \"cqt\"], [\"deep\", \"clp\"], [\"stft\"]]):\n        for row, type in enumerate(types):\n            ch = raw_chroma(audio, sr, type=type)\n            if ch.shape[1] == 12:\n                ch = ch.T\n            librosa.display.specshow(ch, y_axis=\"chroma\", x_axis=\"time\", ax=ax[row, col])\n            ax[row, col].set(title=type)\n            ax[row, col].label_outer()\n    plt.tight_layout()\n    plt.show()\n"
  },
  {
    "path": "augment.py",
    "content": "import math\r\n\r\nimport torch\r\nfrom torch.nn import functional as F\r\n\r\nfrom op import upfirdn2d\r\n\r\n\r\nSYM6 = (\r\n    0.015404109327027373,\r\n    0.0034907120842174702,\r\n    -0.11799011114819057,\r\n    -0.048311742585633,\r\n    0.4910559419267466,\r\n    0.787641141030194,\r\n    0.3379294217276218,\r\n    -0.07263752278646252,\r\n    -0.021060292512300564,\r\n    0.04472490177066578,\r\n    0.0017677118642428036,\r\n    -0.007800708325034148,\r\n)\r\n\r\n\r\ndef translate_mat(t_x, t_y):\r\n    batch = t_x.shape[0]\r\n\r\n    mat = torch.eye(3).unsqueeze(0).repeat(batch, 1, 1)\r\n    translate = torch.stack((t_x, t_y), 1)\r\n    mat[:, :2, 2] = translate\r\n\r\n    return mat\r\n\r\n\r\ndef rotate_mat(theta):\r\n    batch = theta.shape[0]\r\n\r\n    mat = torch.eye(3).unsqueeze(0).repeat(batch, 1, 1)\r\n    sin_t = torch.sin(theta)\r\n    cos_t = torch.cos(theta)\r\n    rot = torch.stack((cos_t, -sin_t, sin_t, cos_t), 1).view(batch, 2, 2)\r\n    mat[:, :2, :2] = rot\r\n\r\n    return mat\r\n\r\n\r\ndef scale_mat(s_x, s_y):\r\n    batch = s_x.shape[0]\r\n\r\n    mat = torch.eye(3).unsqueeze(0).repeat(batch, 1, 1)\r\n    mat[:, 0, 0] = s_x\r\n    mat[:, 1, 1] = s_y\r\n\r\n    return mat\r\n\r\n\r\ndef translate3d_mat(t_x, t_y, t_z):\r\n    batch = t_x.shape[0]\r\n\r\n    mat = torch.eye(4).unsqueeze(0).repeat(batch, 1, 1)\r\n    translate = torch.stack((t_x, t_y, t_z), 1)\r\n    mat[:, :3, 3] = translate\r\n\r\n    return mat\r\n\r\n\r\ndef rotate3d_mat(axis, theta):\r\n    batch = theta.shape[0]\r\n\r\n    u_x, u_y, u_z = axis\r\n\r\n    eye = torch.eye(3).unsqueeze(0)\r\n    cross = torch.tensor([(0, -u_z, u_y), (u_z, 0, -u_x), (-u_y, u_x, 0)]).unsqueeze(0)\r\n    outer = torch.tensor(axis)\r\n    outer = (outer.unsqueeze(1) * outer).unsqueeze(0)\r\n\r\n    sin_t = torch.sin(theta).view(-1, 1, 1)\r\n    cos_t = torch.cos(theta).view(-1, 1, 1)\r\n\r\n    rot = cos_t * eye + sin_t * cross + (1 - cos_t) * outer\r\n\r\n    eye_4 = torch.eye(4).unsqueeze(0).repeat(batch, 1, 1)\r\n    eye_4[:, :3, :3] = rot\r\n\r\n    return eye_4\r\n\r\n\r\ndef scale3d_mat(s_x, s_y, s_z):\r\n    batch = s_x.shape[0]\r\n\r\n    mat = torch.eye(4).unsqueeze(0).repeat(batch, 1, 1)\r\n    mat[:, 0, 0] = s_x\r\n    mat[:, 1, 1] = s_y\r\n    mat[:, 2, 2] = s_z\r\n\r\n    return mat\r\n\r\n\r\ndef luma_flip_mat(axis, i):\r\n    batch = i.shape[0]\r\n\r\n    eye = torch.eye(4).unsqueeze(0).repeat(batch, 1, 1)\r\n    axis = torch.tensor(axis + (0,))\r\n    flip = 2 * torch.ger(axis, axis) * i.view(-1, 1, 1)\r\n\r\n    return eye - flip\r\n\r\n\r\ndef saturation_mat(axis, i):\r\n    batch = i.shape[0]\r\n\r\n    eye = torch.eye(4).unsqueeze(0).repeat(batch, 1, 1)\r\n    axis = torch.tensor(axis + (0,))\r\n    axis = torch.ger(axis, axis)\r\n    saturate = axis + (eye - axis) * i.view(-1, 1, 1)\r\n\r\n    return saturate\r\n\r\n\r\ndef lognormal_sample(size, mean=0, std=1):\r\n    return torch.empty(size).log_normal_(mean=mean, std=std)\r\n\r\n\r\ndef category_sample(size, categories):\r\n    category = torch.tensor(categories)\r\n    sample = torch.randint(high=len(categories), size=(size,))\r\n\r\n    return category[sample]\r\n\r\n\r\ndef uniform_sample(size, low, high):\r\n    return torch.empty(size).uniform_(low, high)\r\n\r\n\r\ndef normal_sample(size, mean=0, std=1):\r\n    return torch.empty(size).normal_(mean, std)\r\n\r\n\r\ndef bernoulli_sample(size, p):\r\n    return torch.empty(size).bernoulli_(p)\r\n\r\n\r\ndef random_mat_apply(p, transform, prev, eye):\r\n    size = transform.shape[0]\r\n    select = bernoulli_sample(size, p).view(size, 1, 1)\r\n    select_transform = select * transform + (1 - select) * eye\r\n\r\n    return select_transform @ prev\r\n\r\n\r\ndef sample_affine(p, size, height, width):\r\n    G = torch.eye(3).unsqueeze(0).repeat(size, 1, 1)\r\n    eye = G\r\n\r\n    # flip\r\n    param = category_sample(size, (0, 1))\r\n    Gc = scale_mat(1 - 2.0 * param, torch.ones(size))\r\n    G = random_mat_apply(p, Gc, G, eye)\r\n    # print('flip', G, scale_mat(1 - 2.0 * param, torch.ones(size)), sep='\\n')\r\n\r\n    # 90 rotate\r\n    param = category_sample(size, (0, 3))\r\n    Gc = rotate_mat(-math.pi / 2 * param)\r\n    G = random_mat_apply(p, Gc, G, eye)\r\n    # print('90 rotate', G, rotate_mat(-math.pi / 2 * param), sep='\\n')\r\n\r\n    # integer translate\r\n    param = uniform_sample(size, -0.125, 0.125)\r\n    param_height = torch.round(param * height) / height\r\n    param_width = torch.round(param * width) / width\r\n    Gc = translate_mat(param_width, param_height)\r\n    G = random_mat_apply(p, Gc, G, eye)\r\n    # print('integer translate', G, translate_mat(param_width, param_height), sep='\\n')\r\n\r\n    # isotropic scale\r\n    param = lognormal_sample(size, std=0.2 * math.log(2))\r\n    Gc = scale_mat(param, param)\r\n    G = random_mat_apply(p, Gc, G, eye)\r\n    # print('isotropic scale', G, scale_mat(param, param), sep='\\n')\r\n\r\n    p_rot = 1 - math.sqrt(1 - p)\r\n\r\n    # pre-rotate\r\n    param = uniform_sample(size, -math.pi, math.pi)\r\n    Gc = rotate_mat(-param)\r\n    G = random_mat_apply(p_rot, Gc, G, eye)\r\n    # print('pre-rotate', G, rotate_mat(-param), sep='\\n')\r\n\r\n    # anisotropic scale\r\n    param = lognormal_sample(size, std=0.2 * math.log(2))\r\n    Gc = scale_mat(param, 1 / param)\r\n    G = random_mat_apply(p, Gc, G, eye)\r\n    # print('anisotropic scale', G, scale_mat(param, 1 / param), sep='\\n')\r\n\r\n    # post-rotate\r\n    param = uniform_sample(size, -math.pi, math.pi)\r\n    Gc = rotate_mat(-param)\r\n    G = random_mat_apply(p_rot, Gc, G, eye)\r\n    # print('post-rotate', G, rotate_mat(-param), sep='\\n')\r\n\r\n    # fractional translate\r\n    param = normal_sample(size, std=0.125)\r\n    Gc = translate_mat(param, param)\r\n    G = random_mat_apply(p, Gc, G, eye)\r\n    # print('fractional translate', G, translate_mat(param, param), sep='\\n')\r\n\r\n    return G\r\n\r\n\r\ndef sample_color(p, size):\r\n    C = torch.eye(4).unsqueeze(0).repeat(size, 1, 1)\r\n    eye = C\r\n    axis_val = 1 / math.sqrt(3)\r\n    axis = (axis_val, axis_val, axis_val)\r\n\r\n    # brightness\r\n    param = normal_sample(size, std=0.2)\r\n    Cc = translate3d_mat(param, param, param)\r\n    C = random_mat_apply(p, Cc, C, eye)\r\n\r\n    # contrast\r\n    param = lognormal_sample(size, std=0.5 * math.log(2))\r\n    Cc = scale3d_mat(param, param, param)\r\n    C = random_mat_apply(p, Cc, C, eye)\r\n\r\n    # luma flip\r\n    param = category_sample(size, (0, 1))\r\n    Cc = luma_flip_mat(axis, param)\r\n    C = random_mat_apply(p, Cc, C, eye)\r\n\r\n    # hue rotation\r\n    param = uniform_sample(size, -math.pi, math.pi)\r\n    Cc = rotate3d_mat(axis, param)\r\n    C = random_mat_apply(p, Cc, C, eye)\r\n\r\n    # saturation\r\n    param = lognormal_sample(size, std=1 * math.log(2))\r\n    Cc = saturation_mat(axis, param)\r\n    C = random_mat_apply(p, Cc, C, eye)\r\n\r\n    return C\r\n\r\n\r\ndef make_grid(shape, x0, x1, y0, y1, device):\r\n    n, c, h, w = shape\r\n    grid = torch.empty(n, h, w, 3, device=device)\r\n    grid[:, :, :, 0] = torch.linspace(x0, x1, w, device=device)\r\n    grid[:, :, :, 1] = torch.linspace(y0, y1, h, device=device).unsqueeze(-1)\r\n    grid[:, :, :, 2] = 1\r\n\r\n    return grid\r\n\r\n\r\ndef affine_grid(grid, mat):\r\n    n, h, w, _ = grid.shape\r\n    return (grid.view(n, h * w, 3) @ mat.transpose(1, 2)).view(n, h, w, 2)\r\n\r\n\r\ndef get_padding(G, height, width):\r\n    extreme = G[:, :2, :] @ torch.tensor([(-1.0, -1, 1), (-1, 1, 1), (1, -1, 1), (1, 1, 1)]).t()\r\n\r\n    size = torch.tensor((width, height))\r\n\r\n    pad_low = ((extreme.min(-1).values + 1) * size).clamp(max=0).abs().ceil().max(0).values.to(torch.int64).tolist()\r\n    pad_high = (extreme.max(-1).values * size - size).clamp(min=0).ceil().max(0).values.to(torch.int64).tolist()\r\n\r\n    return pad_low[0], pad_high[0], pad_low[1], pad_high[1]\r\n\r\n\r\ndef try_sample_affine_and_pad(img, p, pad_k, G=None):\r\n    batch, _, height, width = img.shape\r\n\r\n    G_try = G\r\n\r\n    while True:\r\n        if G is None:\r\n            G_try = sample_affine(p, batch, height, width)\r\n\r\n        pad_x1, pad_x2, pad_y1, pad_y2 = get_padding(torch.inverse(G_try), height, width)\r\n\r\n        try:\r\n            img_pad = F.pad(img, (pad_x1 + pad_k, pad_x2 + pad_k, pad_y1 + pad_k, pad_y2 + pad_k), mode=\"reflect\",)\r\n\r\n        except RuntimeError:\r\n            continue\r\n\r\n        break\r\n\r\n    return img_pad, G_try, (pad_x1, pad_x2, pad_y1, pad_y2)\r\n\r\n\r\ndef random_apply_affine(img, p, G=None, antialiasing_kernel=SYM6):\r\n    kernel = antialiasing_kernel\r\n    len_k = len(kernel)\r\n    pad_k = (len_k + 1) // 2\r\n\r\n    kernel = torch.as_tensor(kernel)\r\n    kernel = torch.ger(kernel, kernel).to(img)\r\n    kernel_flip = torch.flip(kernel, (0, 1))\r\n\r\n    img_pad, G, (pad_x1, pad_x2, pad_y1, pad_y2) = try_sample_affine_and_pad(img, p, pad_k, G)\r\n\r\n    p_ux1 = pad_x1\r\n    p_ux2 = pad_x2 + 1\r\n    p_uy1 = pad_y1\r\n    p_uy2 = pad_y2 + 1\r\n    w_p = img_pad.shape[3] - len_k + 1\r\n    h_p = img_pad.shape[2] - len_k + 1\r\n    h_o = img.shape[2]\r\n    w_o = img.shape[3]\r\n\r\n    img_2x = upfirdn2d(img_pad, kernel_flip, up=2)\r\n\r\n    grid = make_grid(\r\n        img_2x.shape,\r\n        -2 * p_ux1 / w_o - 1,\r\n        2 * (w_p - p_ux1) / w_o - 1,\r\n        -2 * p_uy1 / h_o - 1,\r\n        2 * (h_p - p_uy1) / h_o - 1,\r\n        device=img_2x.device,\r\n    ).to(img_2x)\r\n    grid = affine_grid(grid, torch.inverse(G)[:, :2, :].to(img_2x))\r\n    grid = grid * torch.tensor([w_o / w_p, h_o / h_p], device=grid.device) + torch.tensor(\r\n        [(w_o + 2 * p_ux1) / w_p - 1, (h_o + 2 * p_uy1) / h_p - 1], device=grid.device\r\n    )\r\n\r\n    img_affine = F.grid_sample(img_2x, grid, mode=\"bilinear\", align_corners=False, padding_mode=\"zeros\")\r\n\r\n    img_down = upfirdn2d(img_affine, kernel, down=2)\r\n\r\n    end_y = -pad_y2 - 1\r\n    if end_y == 0:\r\n        end_y = img_down.shape[2]\r\n\r\n    end_x = -pad_x2 - 1\r\n    if end_x == 0:\r\n        end_x = img_down.shape[3]\r\n\r\n    img = img_down[:, :, pad_y1:end_y, pad_x1:end_x]\r\n\r\n    return img, G\r\n\r\n\r\ndef apply_color(img, mat):\r\n    batch = img.shape[0]\r\n    img = img.permute(0, 2, 3, 1)\r\n    mat_mul = mat[:, :3, :3].transpose(1, 2).view(batch, 1, 3, 3)\r\n    mat_add = mat[:, :3, 3].view(batch, 1, 1, 3)\r\n    img = img @ mat_mul + mat_add\r\n    img = img.permute(0, 3, 1, 2)\r\n\r\n    return img\r\n\r\n\r\ndef random_apply_color(img, p, C=None):\r\n    if C is None:\r\n        C = sample_color(p, img.shape[0])\r\n\r\n    img = apply_color(img, C.to(img))\r\n\r\n    return img, C\r\n\r\n\r\ndef augment(img, p, transform_matrix=(None, None)):\r\n    img, G = random_apply_affine(img, p, transform_matrix[0])\r\n    img, C = random_apply_color(img, p, transform_matrix[1])\r\n\r\n    return img, (G, C)\r\n"
  },
  {
    "path": "contrastive_learner.py",
    "content": "import copy\nimport random\nfrom functools import wraps\n\nimport torch\nfrom torch import nn\nimport torch.nn.functional as F\n\n\ndef identity(x):\n    return x\n\n\ndef default(val, def_val):\n    return def_val if val is None else val\n\n\ndef flatten(t):\n    return t.reshape(t.shape[0], -1)\n\n\ndef safe_concat(arr, el, dim=0):\n    if arr is None:\n        return el\n    return torch.cat((arr, el), dim=dim)\n\n\ndef singleton(cache_key):\n    def inner_fn(fn):\n        @wraps(fn)\n        def wrapper(self, *args, **kwargs):\n            instance = getattr(self, cache_key)\n            if instance is not None:\n                return instance\n\n            instance = fn(self, *args, **kwargs)\n            setattr(self, cache_key, instance)\n            return instance\n\n        return wrapper\n\n    return inner_fn\n\n\n# losses\n\n\ndef contrastive_loss(queries, keys, temperature=0.1):\n    b, device = queries.shape[0], queries.device\n    logits = queries @ keys.t()\n    logits = logits - logits.max(dim=-1, keepdim=True).values\n    logits /= temperature\n    return F.cross_entropy(logits, torch.arange(b, device=device))\n\n\ndef nt_xent_loss(queries, keys, temperature=0.1):\n    b, device = queries.shape[0], queries.device\n\n    n = b * 2\n    projs = torch.cat((queries, keys))\n    logits = projs @ projs.t()\n\n    mask = torch.eye(n, device=device).bool()\n    logits = logits[~mask].reshape(n, n - 1)\n    logits /= temperature\n\n    labels = torch.cat(((torch.arange(b, device=device) + b - 1), torch.arange(b, device=device)), dim=0)\n    loss = F.cross_entropy(logits, labels, reduction=\"sum\")\n    loss /= 2 * (b - 1)\n    return loss\n\n\n# augmentation utils\n\n\nclass RandomApply(nn.Module):\n    def __init__(self, fn, p):\n        super().__init__()\n        self.fn = fn\n        self.p = p\n\n    def forward(self, x):\n        x_out = []\n        for ex in x:\n            if random.random() > self.p:\n                x_out.append(ex[None, :])\n            else:\n                x_out.append(self.fn(ex))\n        return torch.cat(x_out)\n\n\n# exponential moving average\n\n\nclass EMA:\n    def __init__(self, beta):\n        super().__init__()\n        self.beta = beta\n\n    def update_average(self, old, new):\n        if old is None:\n            return new\n        return old * self.beta + (1 - self.beta) * new\n\n\ndef update_moving_average(ema_updater, ma_model, current_model):\n    for current_params, ma_params in zip(current_model.parameters(), ma_model.parameters()):\n        old_weight, up_weight = ma_params.data, current_params.data\n        ma_params.data = ema_updater.update_average(old_weight, up_weight)\n\n\n# hidden layer extractor class\n\n\nclass OutputHiddenLayer(nn.Module):\n    def __init__(self, net, layer=-2):\n        super().__init__()\n        self.net = net\n        self.layer = layer\n\n        self.hidden = None\n        self._register_hook()\n\n    def _find_layer(self):\n        if type(self.layer) == str:\n            modules = dict([*self.net.named_modules()])\n            return modules.get(self.layer, None)\n        elif type(self.layer) == int:\n            children = [*self.net.children()]\n            return children[self.layer]\n        elif type(self.layer) == tuple:\n            children = [*self.net.children()]\n            grand_children = [*children[self.layer[0]].children()]\n            return grand_children[self.layer[1]]\n        return None\n\n    def _register_hook(self):\n        def hook(_, __, output):\n            self.hidden = output\n\n        layer = self._find_layer()\n        assert layer is not None, f\"hidden layer ({self.layer}) not found\"\n        handle = layer.register_forward_hook(hook)\n\n    def forward(self, x):\n        if self.layer == -1:\n            return self.net(x)\n\n        _ = self.net(x)\n        hidden = self.hidden\n        self.hidden = None\n        assert hidden is not None, f\"hidden layer {self.layer} never emitted an output\"\n        return hidden\n\n\nclass ContrastiveLearner(nn.Module):\n    def __init__(\n        self,\n        net,\n        image_size,\n        hidden_layer=-2,\n        project_hidden=True,\n        project_dim=128,\n        use_nt_xent_loss=False,\n        use_bilinear=False,\n        use_momentum=False,\n        momentum_value=0.999,\n        key_encoder=None,\n        temperature=0.1,\n    ):\n        super().__init__()\n        self.net = OutputHiddenLayer(net, layer=hidden_layer)\n\n        self.temperature = temperature\n        self.use_nt_xent_loss = use_nt_xent_loss\n\n        self.project_hidden = project_hidden\n        self.projection = None\n        self.project_dim = project_dim\n\n        self.use_bilinear = use_bilinear\n        self.bilinear_w = None\n\n        self.use_momentum = use_momentum\n        self.ema_updater = EMA(momentum_value)\n        self.key_encoder = key_encoder\n\n        # for accumulating queries and keys across calls\n        self.queries = None\n        self.keys = None\n\n        # send a mock image tensor to instantiate parameters\n        init = torch.randn(1, 3, image_size, image_size, device=\"cuda\")\n        self.forward(init)\n\n    @singleton(\"key_encoder\")\n    def _get_key_encoder(self):\n        key_encoder = copy.deepcopy(self.net)\n        key_encoder._register_hook()\n        return key_encoder\n\n    @singleton(\"bilinear_w\")\n    def _get_bilinear(self, hidden):\n        _, dim = hidden.shape\n        return nn.Parameter(torch.eye(dim, device=device, dtype=dtype)).to(hidden)\n\n    @singleton(\"projection\")\n    def _get_projection_fn(self, hidden):\n        _, dim = hidden.shape\n        return nn.Sequential(\n            nn.Linear(dim, dim, bias=False), nn.LeakyReLU(inplace=True), nn.Linear(dim, self.project_dim, bias=False)\n        ).to(hidden)\n\n    def reset_moving_average(self):\n        assert self.use_momentum, \"must be using momentum method for key encoder\"\n        del self.key_encoder\n        self.key_encoder = None\n\n    def update_moving_average(self):\n        assert self.key_encoder is not None, \"key encoder has not been created yet\"\n        self.key_encoder = update_moving_average(self.ema_updater, self.key_encoder, self.net)\n\n    def calculate_loss(self):\n        assert self.queries is not None and self.keys is not None, \"no queries or keys accumulated\"\n        loss_fn = nt_xent_loss if self.use_nt_xent_loss else contrastive_loss\n        loss = loss_fn(self.queries, self.keys, temperature=self.temperature)\n        self.queries = self.keys = None\n        return loss\n\n    def forward(self, x, aug_x, accumulate=False):\n        b, c, h, w, device = *x.shape, x.device\n\n        queries = self.net(aug_x)\n\n        key_encoder = self.net if not self.use_momentum else self._get_key_encoder()\n        keys = key_encoder(aug_x)\n\n        if self.use_momentum:\n            keys = keys.detach()\n\n        queries, keys = map(flatten, (queries, keys))\n\n        if self.use_bilinear:\n            W = self._get_bilinear(keys)\n            keys = (W @ keys.t()).t()\n\n        project_fn = self._get_projection_fn(queries) if self.project_hidden else identity\n        queries, keys = map(project_fn, (queries, keys))\n\n        self.queries = safe_concat(self.queries, queries)\n        self.keys = safe_concat(self.keys, keys)\n\n        return self.calculate_loss() if not accumulate else None\n"
  },
  {
    "path": "convert_weight.py",
    "content": "import argparse\nimport math\nimport os\nimport pickle\nimport sys\n\nimport numpy as np\nimport torch\nfrom torchvision import utils\n\nfrom model import Discriminator, Generator\n\n\ndef convert_modconv(vars, source_name, target_name, flip=False):\n    weight = vars[source_name + \"/weight\"].value().eval()\n    mod_weight = vars[source_name + \"/mod_weight\"].value().eval()\n    mod_bias = vars[source_name + \"/mod_bias\"].value().eval()\n    noise = vars[source_name + \"/noise_strength\"].value().eval()\n    bias = vars[source_name + \"/bias\"].value().eval()\n\n    dic = {\n        \"conv.weight\": np.expand_dims(weight.transpose((3, 2, 0, 1)), 0),\n        \"conv.modulation.weight\": mod_weight.transpose((1, 0)),\n        \"conv.modulation.bias\": mod_bias + 1,\n        \"noise.weight\": np.array([noise]),\n        \"activate.bias\": bias,\n    }\n\n    dic_torch = {}\n\n    for k, v in dic.items():\n        dic_torch[target_name + \".\" + k] = torch.from_numpy(v)\n\n    if flip:\n        dic_torch[target_name + \".conv.weight\"] = torch.flip(dic_torch[target_name + \".conv.weight\"], [3, 4])\n\n    return dic_torch\n\n\ndef convert_conv(vars, source_name, target_name, bias=True, start=0):\n    weight = vars[source_name + \"/weight\"].value().eval()\n\n    dic = {\"weight\": weight.transpose((3, 2, 0, 1))}\n\n    if bias:\n        dic[\"bias\"] = vars[source_name + \"/bias\"].value().eval()\n\n    dic_torch = {}\n\n    dic_torch[target_name + f\".{start}.weight\"] = torch.from_numpy(dic[\"weight\"])\n\n    if bias:\n        dic_torch[target_name + f\".{start + 1}.bias\"] = torch.from_numpy(dic[\"bias\"])\n\n    return dic_torch\n\n\ndef convert_torgb(vars, source_name, target_name):\n    weight = vars[source_name + \"/weight\"].value().eval()\n    mod_weight = vars[source_name + \"/mod_weight\"].value().eval()\n    mod_bias = vars[source_name + \"/mod_bias\"].value().eval()\n    bias = vars[source_name + \"/bias\"].value().eval()\n\n    dic = {\n        \"conv.weight\": np.expand_dims(weight.transpose((3, 2, 0, 1)), 0),\n        \"conv.modulation.weight\": mod_weight.transpose((1, 0)),\n        \"conv.modulation.bias\": mod_bias + 1,\n        \"bias\": bias.reshape((1, 3, 1, 1)),\n    }\n\n    dic_torch = {}\n\n    for k, v in dic.items():\n        dic_torch[target_name + \".\" + k] = torch.from_numpy(v)\n\n    return dic_torch\n\n\ndef convert_dense(vars, source_name, target_name):\n    weight = vars[source_name + \"/weight\"].value().eval()\n    bias = vars[source_name + \"/bias\"].value().eval()\n\n    dic = {\"weight\": weight.transpose((1, 0)), \"bias\": bias}\n\n    dic_torch = {}\n\n    for k, v in dic.items():\n        dic_torch[target_name + \".\" + k] = torch.from_numpy(v)\n\n    return dic_torch\n\n\ndef update(state_dict, new):\n    for k, v in new.items():\n        if k not in state_dict:\n            raise KeyError(k + \" is not found\")\n\n        if v.shape != state_dict[k].shape:\n            raise ValueError(f\"Shape mismatch: {v.shape} vs {state_dict[k].shape}\")\n\n        state_dict[k] = v\n\n\ndef discriminator_fill_statedict(statedict, vars, size):\n    log_size = int(math.log(size, 2))\n\n    update(statedict, convert_conv(vars, f\"{size}x{size}/FromRGB\", \"convs.0\"))\n\n    conv_i = 1\n\n    for i in range(log_size - 2, 0, -1):\n        reso = 4 * 2 ** i\n        update(\n            statedict, convert_conv(vars, f\"{reso}x{reso}/Conv0\", f\"convs.{conv_i}.conv1\"),\n        )\n        update(\n            statedict, convert_conv(vars, f\"{reso}x{reso}/Conv1_down\", f\"convs.{conv_i}.conv2\", start=1),\n        )\n        update(\n            statedict, convert_conv(vars, f\"{reso}x{reso}/Skip\", f\"convs.{conv_i}.skip\", start=1, bias=False),\n        )\n        conv_i += 1\n\n    update(statedict, convert_conv(vars, f\"4x4/Conv\", \"final_conv\"))\n    update(statedict, convert_dense(vars, f\"4x4/Dense0\", \"final_linear.0\"))\n    update(statedict, convert_dense(vars, f\"Output\", \"final_linear.1\"))\n\n    return statedict\n\n\ndef fill_statedict(state_dict, vars, size):\n    log_size = int(math.log(size, 2))\n\n    for i in range(8):\n        update(state_dict, convert_dense(vars, f\"G_mapping/Dense{i}\", f\"style.{i + 1}\"))\n\n    update(\n        state_dict, {\"input.input\": torch.from_numpy(vars[\"G_synthesis/4x4/Const/const\"].value().eval())},\n    )\n\n    update(state_dict, convert_torgb(vars, \"G_synthesis/4x4/ToRGB\", \"to_rgb1\"))\n\n    for i in range(log_size - 2):\n        reso = 4 * 2 ** (i + 1)\n        update(\n            state_dict, convert_torgb(vars, f\"G_synthesis/{reso}x{reso}/ToRGB\", f\"to_rgbs.{i}\"),\n        )\n\n    update(state_dict, convert_modconv(vars, \"G_synthesis/4x4/Conv\", \"conv1\"))\n\n    conv_i = 0\n\n    for i in range(log_size - 2):\n        reso = 4 * 2 ** (i + 1)\n        update(\n            state_dict, convert_modconv(vars, f\"G_synthesis/{reso}x{reso}/Conv0_up\", f\"convs.{conv_i}\", flip=True,),\n        )\n        update(\n            state_dict, convert_modconv(vars, f\"G_synthesis/{reso}x{reso}/Conv1\", f\"convs.{conv_i + 1}\"),\n        )\n        conv_i += 2\n\n    for i in range(0, (log_size - 2) * 2 + 1):\n        update(\n            state_dict, {f\"noises.noise_{i}\": torch.from_numpy(vars[f\"G_synthesis/noise{i}\"].value().eval())},\n        )\n\n    return state_dict\n\n\nif __name__ == \"__main__\":\n    device = \"cuda\"\n\n    parser = argparse.ArgumentParser()\n    parser.add_argument(\"--repo\", type=str, required=True)\n    parser.add_argument(\"--gen\", action=\"store_true\")\n    parser.add_argument(\"--disc\", action=\"store_true\")\n    parser.add_argument(\"--channel_multiplier\", type=int, default=2)\n    parser.add_argument(\"path\", metavar=\"PATH\")\n\n    args = parser.parse_args()\n\n    sys.path.append(args.repo)\n\n    import dnnlib\n    from dnnlib import tflib\n\n    tflib.init_tf()\n\n    with open(args.path, \"rb\") as f:\n        generator, discriminator, g_ema = pickle.load(f)\n\n    size = g_ema.output_shape[2]\n\n    g = Generator(size, 512, 8, channel_multiplier=args.channel_multiplier)\n    state_dict = g.state_dict()\n    state_dict = fill_statedict(state_dict, g_ema.vars, size)\n\n    g.load_state_dict(state_dict)\n\n    latent_avg = torch.from_numpy(g_ema.vars[\"dlatent_avg\"].value().eval())\n\n    ckpt = {\"g_ema\": state_dict, \"latent_avg\": latent_avg}\n\n    if args.gen:\n        g_train = Generator(size, 512, 8, channel_multiplier=args.channel_multiplier)\n        g_train_state = g_train.state_dict()\n        g_train_state = fill_statedict(g_train_state, generator.vars, size)\n        ckpt[\"g\"] = g_train_state\n\n    if args.disc:\n        disc = Discriminator(size, channel_multiplier=args.channel_multiplier)\n        d_state = disc.state_dict()\n        d_state = discriminator_fill_statedict(d_state, discriminator.vars, size)\n        ckpt[\"d\"] = d_state\n\n    name = os.path.splitext(os.path.basename(args.path))[0]\n    torch.save(ckpt, name + \".pt\")\n\n    batch_size = {256: 16, 512: 9, 1024: 4}\n    n_sample = batch_size.get(size, 25)\n\n    g = g.to(device)\n\n    z = np.random.RandomState(0).randn(n_sample, 512).astype(\"float32\")\n\n    with torch.no_grad():\n        img_pt, _ = g([torch.from_numpy(z).to(device)], truncation=0.5, truncation_latent=latent_avg.to(device),)\n\n    Gs_kwargs = dnnlib.EasyDict()\n    Gs_kwargs.randomize_noise = False\n    img_tf = g_ema.run(z, None, **Gs_kwargs)\n    img_tf = torch.from_numpy(img_tf).to(device)\n\n    img_diff = ((img_pt + 1) / 2).clamp(0.0, 1.0) - ((img_tf.to(device) + 1) / 2).clamp(0.0, 1.0)\n\n    img_concat = torch.cat((img_tf, img_pt, img_diff), dim=0)\n    utils.save_image(img_concat, name + \".png\", nrow=n_sample, normalize=True, range=(-1, 1))\n"
  },
  {
    "path": "dataset.py",
    "content": "from io import BytesIO\n\nimport lmdb\nimport numpy as np\nfrom PIL import Image\nfrom PIL import Image\nfrom torch.utils.data import Dataset\n\n\nclass MultiResolutionDataset(Dataset):\n    def __init__(self, path, transform, resolution=256):\n        self.env = lmdb.open(path, max_readers=32, readonly=True, lock=False, readahead=False, meminit=False,)\n\n        if not self.env:\n            raise IOError(\"Cannot open lmdb dataset\", path)\n\n        with self.env.begin(write=False) as txn:\n            self.length = int(txn.get(\"length\".encode(\"utf-8\")).decode(\"utf-8\"))\n\n        self.resolution = resolution\n        self.transform = transform\n\n    def __len__(self):\n        return self.length\n\n    def __getitem__(self, index):\n        while True:\n            try:\n                with self.env.begin(write=False) as txn:\n                    key = f\"{self.resolution}-{str(index).zfill(5)}\".encode(\"utf-8\")\n                    img_bytes = txn.get(key)\n\n                buffer = BytesIO(img_bytes)\n                img = Image.open(buffer)\n                break\n            except:\n                print(f\"ERROR loading image {index}\")\n                index = int(np.random.rand() * self.length)\n                print(f\"Trying again with {index}...\")\n        img = self.transform(img)\n\n        return img\n"
  },
  {
    "path": "distributed.py",
    "content": "import pickle\n\nimport torch\nfrom torch import distributed as dist\n\n\ndef get_rank():\n    if not dist.is_available():\n        return 0\n\n    if not dist.is_initialized():\n        return 0\n\n    return dist.get_rank()\n\n\ndef synchronize():\n    if not dist.is_available():\n        return\n\n    if not dist.is_initialized():\n        return\n\n    world_size = dist.get_world_size()\n\n    if world_size == 1:\n        return\n\n    dist.barrier()\n\n\ndef get_world_size():\n    if not dist.is_available():\n        return 1\n\n    if not dist.is_initialized():\n        return 1\n\n    return dist.get_world_size()\n\n\ndef reduce_sum(tensor):\n    if not dist.is_available():\n        return tensor\n\n    if not dist.is_initialized():\n        return tensor\n\n    tensor = tensor.clone()\n    dist.all_reduce(tensor, op=dist.ReduceOp.SUM)\n\n    return tensor\n\n\ndef gather_grad(params):\n    world_size = get_world_size()\n\n    if world_size == 1:\n        return\n\n    for param in params:\n        if param.grad is not None:\n            dist.all_reduce(param.grad.data, op=dist.ReduceOp.SUM)\n            param.grad.data.div_(world_size)\n\n\ndef all_gather(data):\n    world_size = get_world_size()\n\n    if world_size == 1:\n        return [data]\n\n    buffer = pickle.dumps(data)\n    storage = torch.ByteStorage.from_buffer(buffer)\n    tensor = torch.ByteTensor(storage).to(\"cuda\")\n\n    local_size = torch.IntTensor([tensor.numel()]).to(\"cuda\")\n    size_list = [torch.IntTensor([0]).to(\"cuda\") for _ in range(world_size)]\n    dist.all_gather(size_list, local_size)\n    size_list = [int(size.item()) for size in size_list]\n    max_size = max(size_list)\n\n    tensor_list = []\n    for _ in size_list:\n        tensor_list.append(torch.ByteTensor(size=(max_size,)).to(\"cuda\"))\n\n    if local_size != max_size:\n        padding = torch.ByteTensor(size=(max_size - local_size,)).to(\"cuda\")\n        tensor = torch.cat((tensor, padding), 0)\n\n    dist.all_gather(tensor_list, tensor)\n\n    data_list = []\n\n    for size, tensor in zip(size_list, tensor_list):\n        buffer = tensor.cpu().numpy().tobytes()[:size]\n        data_list.append(pickle.loads(buffer))\n\n    return data_list\n\n\ndef reduce_loss_dict(loss_dict):\n    world_size = get_world_size()\n\n    if world_size < 2:\n        return loss_dict\n\n    with torch.no_grad():\n        keys = []\n        losses = []\n\n        for k in sorted(loss_dict.keys()):\n            keys.append(k)\n            losses.append(loss_dict[k])\n\n        losses = torch.stack(losses, 0)\n        dist.reduce(losses, dst=0)\n\n        if dist.get_rank() == 0:\n            losses /= world_size\n\n        reduced_losses = {k: v for k, v in zip(keys, losses)}\n\n    return reduced_losses\n"
  },
  {
    "path": "generate.py",
    "content": "import argparse\nimport torch\nfrom torchvision import utils\nfrom models.stylegan2 import Generator\nfrom tqdm import tqdm\n\n\ndef generate(args, g_ema, device, mean_latent):\n\n    with torch.no_grad():\n        g_ema.eval()\n        for i in tqdm(range(args.pics)):\n            sample_z = torch.randn(args.sample, args.latent, device=device)\n\n            sample, _ = g_ema([sample_z], truncation=args.truncation, truncation_latent=mean_latent)\n\n            utils.save_image(\n                sample, f\"sample/{str(i).zfill(6)}.png\", nrow=1, normalize=True, range=(-1, 1),\n            )\n\n\nif __name__ == \"__main__\":\n    device = \"cuda\"\n\n    parser = argparse.ArgumentParser()\n\n    parser.add_argument(\"--size\", type=int, default=1024)\n    parser.add_argument(\"--sample\", type=int, default=1)\n    parser.add_argument(\"--pics\", type=int, default=20)\n    parser.add_argument(\"--truncation\", type=float, default=1)\n    parser.add_argument(\"--truncation_mean\", type=int, default=4096)\n    parser.add_argument(\"--ckpt\", type=str, default=\"stylegan2-ffhq-config-f.pt\")\n    parser.add_argument(\"--channel_multiplier\", type=int, default=2)\n\n    args = parser.parse_args()\n\n    args.latent = 512\n    args.n_mlp = 8\n\n    g_ema = Generator(args.size, args.latent, args.n_mlp, channel_multiplier=args.channel_multiplier).to(device)\n    checkpoint = torch.load(args.ckpt)\n\n    g_ema.load_state_dict(checkpoint[\"g_ema\"])\n\n    if args.truncation < 1:\n        with torch.no_grad():\n            mean_latent = g_ema.mean_latent(args.truncation_mean)\n    else:\n        mean_latent = None\n\n    generate(args, g_ema, device, mean_latent)\n"
  },
  {
    "path": "generate_audiovisual.py",
    "content": "import argparse\nimport gc\nimport os\nimport random\nimport time\nimport traceback\nimport uuid\nimport warnings\n\nimport librosa as rosa\nimport librosa.display\nimport numpy as np\nimport torch as th\n\nimport audioreactive as ar\nimport generate\nimport render\nfrom models.stylegan1 import G_style\nfrom models.stylegan2 import Generator\n\n\ndef get_noise_range(out_size, generator_resolution, is_stylegan1):\n    \"\"\"Gets the correct number of noise resolutions for a given resolution of StyleGAN 1 or 2\"\"\"\n    log_max_res = int(np.log2(out_size))\n    log_min_res = 2 + (log_max_res - int(np.log2(generator_resolution)))\n    if is_stylegan1:\n        range_min = log_min_res\n        range_max = log_max_res + 1\n        side_fn = lambda x: x\n    else:\n        range_min = 2 * log_min_res + 1\n        range_max = 2 * (log_max_res + 1)\n        side_fn = lambda x: int(x / 2)\n    return range_min, range_max, side_fn\n\n\ndef load_generator(\n    ckpt, is_stylegan1, G_res, out_size, noconst, latent_dim, n_mlp, channel_multiplier, dataparallel, base_res_factor\n):\n    \"\"\"Loads a StyleGAN 1 or 2 generator\"\"\"\n    if is_stylegan1:\n        generator = G_style(output_size=out_size, checkpoint=ckpt).cuda()\n    else:\n        generator = Generator(\n            G_res,\n            latent_dim,\n            n_mlp,\n            channel_multiplier=channel_multiplier,\n            constant_input=not noconst,\n            checkpoint=ckpt,\n            output_size=out_size,\n            base_res_factor=base_res_factor,\n        ).cuda()\n    if dataparallel:\n        generator = th.nn.DataParallel(generator)\n    return generator\n\n\ndef generate(\n    ckpt,\n    audio_file,\n    initialize=None,\n    get_latents=None,\n    get_noise=None,\n    get_bends=None,\n    get_rewrites=None,\n    get_truncation=None,\n    output_dir=\"./output\",\n    audioreactive_file=\"audioreactive/examples/default.py\",\n    offset=0,\n    duration=-1,\n    latent_file=None,\n    shuffle_latents=False,\n    G_res=1024,\n    out_size=1024,\n    fps=30,\n    latent_count=12,\n    batch=8,\n    dataparallel=False,\n    truncation=1.0,\n    stylegan1=False,\n    noconst=False,\n    latent_dim=512,\n    n_mlp=8,\n    channel_multiplier=2,\n    randomize_noise=False,\n    ffmpeg_preset=\"slow\",\n    base_res_factor=1,\n    output_file=None,\n    args=None,\n):\n    # if args is empty (i.e. generate() called directly instead of through __main__)\n    # create args Namespace with all locally available variables\n    if args is None:\n        kwargs = locals()\n        args = argparse.Namespace()\n        for k, v in kwargs.items():\n            setattr(args, k, v)\n\n    # ensures smoothing is independent of frame rate\n    ar.set_SMF(args.fps / 30)\n\n    time_taken = time.time()\n    th.set_grad_enabled(False)\n\n    audio, sr, duration = ar.load_audio(audio_file, offset, duration)\n\n    args.audio = audio\n    args.sr = sr\n\n    n_frames = int(round(duration * fps))\n    args.duration = duration\n    args.n_frames = n_frames\n\n    if initialize is not None:\n        args = initialize(args)\n\n    # ====================================================================================\n    # =========================== generate audiovisual latents ===========================\n    # ====================================================================================\n    print(\"\\ngenerating latents...\")\n    if get_latents is None:\n        from audioreactive.default import get_latents\n\n    if latent_file is not None:\n        latent_selection = ar.load_latents(latent_file)\n    else:\n        latent_selection = ar.generate_latents(\n            args.latent_count, ckpt, G_res, noconst, latent_dim, n_mlp, channel_multiplier\n        )\n    if shuffle_latents:\n        random_indices = random.sample(range(len(latent_selection)), len(latent_selection))\n        latent_selection = latent_selection[random_indices]\n    np.save(\"workspace/last-latents.npy\", latent_selection.numpy())\n\n    latents = get_latents(selection=latent_selection, args=args).cpu()\n\n    print(f\"{list(latents.shape)} amplitude={latents.std()}\\n\")\n\n    # ====================================================================================\n    # ============================ generate audiovisual noise ============================\n    # ====================================================================================\n    print(\"generating noise...\")\n    if get_noise is None:\n        from audioreactive.default import get_noise\n\n    noise = []\n    range_min, range_max, exponent = get_noise_range(out_size, G_res, stylegan1)\n    for scale in range(range_min, range_max):\n        h = (2 if out_size == 1080 else 1) * 2 ** exponent(scale)\n        w = (2 if out_size == 1920 else 1) * 2 ** exponent(scale)\n\n        noise.append(get_noise(height=h, width=w, scale=scale - range_min, num_scales=range_max - range_min, args=args))\n\n        if noise[-1] is not None:\n            print(list(noise[-1].shape), f\"amplitude={noise[-1].std()}\")\n        gc.collect()\n        th.cuda.empty_cache()\n    print()\n\n    # ====================================================================================\n    # ================ generate audiovisual network bending manipulations ================\n    # ====================================================================================\n    if get_bends is not None:\n        print(\"generating network bends...\")\n        bends = get_bends(args=args)\n    else:\n        bends = []\n\n    # ====================================================================================\n    # ================ generate audiovisual model rewriting manipulations ================\n    # ====================================================================================\n    if get_rewrites is not None:\n        print(\"generating model rewrites...\")\n        rewrites = get_rewrites(args=args)\n    else:\n        rewrites = {}\n\n    # ====================================================================================\n    # ========================== generate audiovisual truncation =========================\n    # ====================================================================================\n    if get_truncation is not None:\n        print(\"generating truncation...\")\n        truncation = get_truncation(args=args)\n    else:\n        truncation = float(truncation)\n\n    # ====================================================================================\n    # ==== render the given (latent, noise, bends, rewrites, truncation) interpolation ===\n    # ====================================================================================\n    gc.collect()\n    th.cuda.empty_cache()\n\n    generator = load_generator(\n        ckpt=ckpt,\n        is_stylegan1=stylegan1,\n        G_res=G_res,\n        out_size=out_size,\n        noconst=noconst,\n        latent_dim=latent_dim,\n        n_mlp=n_mlp,\n        channel_multiplier=channel_multiplier,\n        dataparallel=dataparallel,\n        base_res_factor=base_res_factor,\n    )\n\n    print(f\"\\npreprocessing took {time.time() - time_taken:.2f}s\\n\")\n\n    print(f\"rendering {n_frames} frames...\")\n    if output_file is None:\n        checkpoint_title = ckpt.split(\"/\")[-1].split(\".\")[0].lower()\n        track_title = audio_file.split(\"/\")[-1].split(\".\")[0].lower()\n        output_file = f\"{output_dir}/{track_title}_{checkpoint_title}_{uuid.uuid4().hex[:8]}.mp4\"\n    render.render(\n        generator=generator,\n        latents=latents,\n        noise=noise,\n        audio_file=audio_file,\n        offset=offset,\n        duration=duration,\n        batch_size=batch,\n        truncation=truncation,\n        bends=bends,\n        rewrites=rewrites,\n        out_size=out_size,\n        output_file=output_file,\n        randomize_noise=randomize_noise,\n        ffmpeg_preset=ffmpeg_preset,\n    )\n\n    print(f\"\\ntotal time taken: {(time.time() - time_taken)/60:.2f} minutes\")\n\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser()\n    parser.add_argument(\"--ckpt\", type=str)\n    parser.add_argument(\"--audio_file\", type=str)\n    parser.add_argument(\"--audioreactive_file\", type=str, default=\"audioreactive/examples/default.py\")\n    parser.add_argument(\"--output_dir\", type=str, default=\"./output\")\n    parser.add_argument(\"--offset\", type=float, default=0)\n    parser.add_argument(\"--duration\", type=float, default=-1, help=\"length of rendered video in seconds\")\n    parser.add_argument(\"--latent_file\", type=str, default=None)\n    parser.add_argument(\"--shuffle_latents\", action=\"store_true\")\n    parser.add_argument(\"--G_res\", type=int, default=1024)\n    parser.add_argument(\"--out_size\", type=int, default=1024, help=\"rendered video size. Options: 512, 1024, 1920\")\n    parser.add_argument(\"--fps\", type=int, default=30)\n    parser.add_argument(\"--latent_count\", type=int, default=12)\n    parser.add_argument(\"--batch\", type=int, default=8)\n    parser.add_argument(\"--dataparallel\", action=\"store_true\")\n    parser.add_argument(\"--truncation\", type=float, default=1.0)\n    parser.add_argument(\"--stylegan1\", action=\"store_true\")\n    parser.add_argument(\"--noconst\", action=\"store_true\")\n    parser.add_argument(\"--latent_dim\", type=int, default=512)\n    parser.add_argument(\"--n_mlp\", type=int, default=8)\n    parser.add_argument(\"--channel_multiplier\", type=int, default=2)\n    parser.add_argument(\"--randomize_noise\", action=\"store_true\")\n    parser.add_argument(\"--base_res_factor\", type=float, default=1)\n    parser.add_argument(\"--ffmpeg_preset\", type=str, default=\"slow\")\n    parser.add_argument(\"--output_file\", type=str, default=None)\n    args = parser.parse_args()\n\n    # ensure output_dir exists\n    os.makedirs(args.output_dir, exist_ok=True)\n\n    # transform file path to python module string\n    modnames = args.audioreactive_file.replace(\".py\", \"\").replace(\"/\", \".\").split(\".\")\n\n    # try to load each of the standard functions from the specified file\n    func_names = [\"initialize\", \"get_latents\", \"get_noise\", \"get_bends\", \"get_rewrites\", \"get_truncation\"]\n    funcs = {}\n    for func in func_names:\n        try:\n            file = __import__(\".\".join(modnames[:-1]), fromlist=[modnames[-1]]).__dict__[modnames[-1]]\n            funcs[func] = getattr(file, func)\n        except AttributeError as error:\n            print(f\"No '{func}' function found in --audioreactive_file, using default...\")\n            funcs[func] = None\n        except:\n            if funcs.get(func, \"error\") == \"error\":\n                print(\"Error while loading --audioreactive_file...\")\n                traceback.print_exc()\n                exit(1)\n\n    # override with args from the OVERRIDE dict in the specified file\n    arg_dict = vars(args).copy()\n    try:\n        file = __import__(\".\".join(modnames[:-1]), fromlist=[modnames[-1]]).__dict__[modnames[-1]]\n        for arg, val in getattr(file, \"OVERRIDE\").items():\n            arg_dict[arg] = val\n            setattr(args, arg, val)\n    except:\n        pass  # no overrides, just continue\n\n    ckpt = arg_dict.pop(\"ckpt\", None)\n    audio_file = arg_dict.pop(\"audio_file\", None)\n\n    # splat kwargs to function call\n    # (generate() has all kwarg defaults specified again to make it amenable to ipynb usage)\n    generate(ckpt=ckpt, audio_file=audio_file, **funcs, **arg_dict, args=args)\n"
  },
  {
    "path": "generate_video.py",
    "content": "import argparse\nimport uuid\n\nimport numpy as np\nimport torch as th\nimport torch.multiprocessing as mp\nimport torch.nn.functional as F\n\nfrom models.stylegan1 import G_style\nfrom models.stylegan2 import Generator\nfrom render import render\n\n\ndef gaussian_filter(x, sigma):\n    dim = len(x.shape)\n    if dim != 3 and dim != 4:\n        raise Exception(\"Only 3- or 4-dimensional tensors are supported.\")\n\n    radius = sigma * 4\n    channels = x.shape[1]\n\n    kernel = th.arange(-radius, radius + 1, dtype=th.float32, device=\"cuda\")\n    kernel = th.exp(-0.5 / sigma ** 2 * kernel ** 2)\n    kernel = kernel / kernel.sum()\n    kernel = kernel.view(1, 1, len(kernel)).repeat(channels, 1, 1)\n\n    if dim == 4:\n        t, c, h, w = x.shape\n        x = x.view(t, c, h * w)\n    x = x.transpose(0, 2)\n\n    x = F.pad(x, (radius, radius), mode=\"circular\")\n    x = F.conv1d(x, weight=kernel, groups=channels)\n\n    x = x.transpose(0, 2)\n    if dim == 4:\n        x = x.view(t, c, h, w)\n\n    return x\n\n\ndef slerp(val, low, high):\n    omega = np.arccos(np.clip(np.dot(low / np.linalg.norm(low), high / np.linalg.norm(high)), -1, 1))\n    so = np.sin(omega)\n    if so == 0:\n        return (1.0 - val) * low + val * high  # L'Hopital's rule/LERP\n    return np.sin((1.0 - val) * omega) / so * low + np.sin(val * omega) / so * high\n\n\ndef lerp(val, low, high):\n    return (1 - val) * low + val * high\n\n\ndef interpolant(t):\n    return t * t * t * (t * (t * 6 - 15) + 10)\n\n\ndef perlin_noise(shape, res, tileable=(True, False, False), interpolant=interpolant):\n    \"\"\"Generate a 3D tensor of perlin noise.\n    Args:\n        shape: The shape of the generated tensor (tuple of three ints).\n            This must be a multiple of res.\n        res: The number of periods of noise to generate along each\n            axis (tuple of three ints). Note shape must be a multiple\n            of res.\n        tileable: If the noise should be tileable along each axis\n            (tuple of three bools). Defaults to (False, False, False).\n        interpolant: The interpolation function, defaults to\n            t*t*t*(t*(t*6 - 15) + 10).\n    Returns:\n        A tensor of shape shape with the generated noise.\n    Raises:\n        ValueError: If shape is not a multiple of res.\n    \"\"\"\n    delta = (res[0] / shape[0], res[1] / shape[1], res[2] / shape[2])\n    d = (shape[0] // res[0], shape[1] // res[1], shape[2] // res[2])\n    grid = np.mgrid[0 : res[0] : delta[0], 0 : res[1] : delta[1], 0 : res[2] : delta[2]]\n    # print(np.mgrid[0 : res[0] : delta[0]])\n    # print(0, res[0], delta[0])\n    # print(th.linspace(0, res[0], delta[0]))\n    # grid = th.meshgrid(\n    #     th.linspace(0, res[0], delta[0]), th.linspace(0, res[1], delta[1]), th.linspace(0, res[1], delta[1])\n    # ).cuda()\n    grid = grid.transpose(1, 2, 3, 0) % 1\n    grid = th.from_numpy(grid).cuda()\n    # Gradients\n    theta = 2 * np.pi * np.random.rand(res[0] + 1, res[1] + 1, res[2] + 1)\n    phi = 2 * np.pi * np.random.rand(res[0] + 1, res[1] + 1, res[2] + 1)\n    gradients = np.stack((np.sin(phi) * np.cos(theta), np.sin(phi) * np.sin(theta), np.cos(phi)), axis=3)\n    if tileable[0]:\n        gradients[-1, :, :] = gradients[0, :, :]\n    if tileable[1]:\n        gradients[:, -1, :] = gradients[:, 0, :]\n    if tileable[2]:\n        gradients[:, :, -1] = gradients[:, :, 0]\n    gradients = gradients.repeat(d[0], 0).repeat(d[1], 1).repeat(d[2], 2)\n    gradients = th.from_numpy(gradients).cuda()\n    g000 = gradients[: -d[0], : -d[1], : -d[2]]\n    g100 = gradients[d[0] :, : -d[1], : -d[2]]\n    g010 = gradients[: -d[0], d[1] :, : -d[2]]\n    g110 = gradients[d[0] :, d[1] :, : -d[2]]\n    g001 = gradients[: -d[0], : -d[1], d[2] :]\n    g101 = gradients[d[0] :, : -d[1], d[2] :]\n    g011 = gradients[: -d[0], d[1] :, d[2] :]\n    g111 = gradients[d[0] :, d[1] :, d[2] :]\n    # Ramps\n    n000 = th.sum(th.stack((grid[:, :, :, 0], grid[:, :, :, 1], grid[:, :, :, 2]), axis=3) * g000, 3)\n    n100 = th.sum(th.stack((grid[:, :, :, 0] - 1, grid[:, :, :, 1], grid[:, :, :, 2]), axis=3) * g100, 3)\n    n010 = th.sum(th.stack((grid[:, :, :, 0], grid[:, :, :, 1] - 1, grid[:, :, :, 2]), axis=3) * g010, 3)\n    n110 = th.sum(th.stack((grid[:, :, :, 0] - 1, grid[:, :, :, 1] - 1, grid[:, :, :, 2]), axis=3) * g110, 3)\n    n001 = th.sum(th.stack((grid[:, :, :, 0], grid[:, :, :, 1], grid[:, :, :, 2] - 1), axis=3) * g001, 3)\n    n101 = th.sum(th.stack((grid[:, :, :, 0] - 1, grid[:, :, :, 1], grid[:, :, :, 2] - 1), axis=3) * g101, 3)\n    n011 = th.sum(th.stack((grid[:, :, :, 0], grid[:, :, :, 1] - 1, grid[:, :, :, 2] - 1), axis=3) * g011, 3)\n    n111 = th.sum(th.stack((grid[:, :, :, 0] - 1, grid[:, :, :, 1] - 1, grid[:, :, :, 2] - 1), axis=3) * g111, 3)\n    # Interpolation\n    t = interpolant(grid)\n    n00 = n000 * (1 - t[:, :, :, 0]) + t[:, :, :, 0] * n100\n    n10 = n010 * (1 - t[:, :, :, 0]) + t[:, :, :, 0] * n110\n    n01 = n001 * (1 - t[:, :, :, 0]) + t[:, :, :, 0] * n101\n    n11 = n011 * (1 - t[:, :, :, 0]) + t[:, :, :, 0] * n111\n    n0 = (1 - t[:, :, :, 1]) * n00 + t[:, :, :, 1] * n10\n    n1 = (1 - t[:, :, :, 1]) * n01 + t[:, :, :, 1] * n11\n    return (1 - t[:, :, :, 2]) * n0 + t[:, :, :, 2] * n1\n\n\ndef spline_loops(base_latent_selection, loop_starting_latents, n_frames, num_loops, smoothing, s=True):\n    from scipy import interpolate\n\n    base_latent_selection = np.concatenate([base_latent_selection, base_latent_selection[[0]]])\n\n    x = np.linspace(0, 1, n_frames // max(1, num_loops))\n    base_latents = np.zeros((len(x), *base_latent_selection.shape[1:]))\n    for lay in range(base_latent_selection.shape[1]):\n        for lat in range(base_latent_selection.shape[2]):\n            tck = interpolate.splrep(\n                np.linspace(0, 1, base_latent_selection.shape[0], dtype=np.float32), base_latent_selection[:, lay, lat],\n            )\n            base_latents[:, lay, lat] = interpolate.splev(x, tck)\n\n    base_latents = th.cat([th.from_numpy(base_latents)] * int(n_frames / len(base_latents)), axis=0).float()\n\n    return base_latents\n\n\ndef get_latent_loops(base_latent_selection, loop_starting_latents, n_frames, num_loops, smoothing, s=True):\n    base_latents = []\n\n    for n in range(len(base_latent_selection)):\n        for val in np.linspace(0.0, 1.0, int(n_frames // max(1, num_loops) // len(base_latent_selection))):\n            base_latents.append(\n                (slerp if s else lerp)(\n                    val,\n                    base_latent_selection[(n + loop_starting_latents) % len(base_latent_selection)][0].cpu(),\n                    base_latent_selection[(n + loop_starting_latents + 1) % len(base_latent_selection)][0].cpu(),\n                )\n            )\n\n    base_latents = th.stack(base_latents, axis=0).cuda()\n    base_latents = th.cat([base_latents] * int(n_frames / len(base_latents)), axis=0)\n    base_latents = th.stack([base_latents] * base_latent_selection.shape[1], axis=1)\n\n    base_latents = gaussian_filter(base_latents, smoothing)\n\n    return base_latents\n\n\nif \"main\" in __name__:\n    parser = argparse.ArgumentParser()\n\n    parser.add_argument(\"--ckpt\", type=str)\n    parser.add_argument(\"--G_res\", type=int, default=1024)\n    parser.add_argument(\"--out_size\", type=int, default=1024)\n    parser.add_argument(\"--batch\", type=int, default=12)\n    parser.add_argument(\"--n_frames\", type=int, default=24 * 30)\n    parser.add_argument(\"--duration\", type=int, default=24)\n    parser.add_argument(\"--const\", type=bool, default=False)\n    parser.add_argument(\"--channel_multiplier\", type=int, default=2)\n    parser.add_argument(\"--truncation\", type=int, default=0.7)\n    parser.add_argument(\"--stylegan1\", type=bool, default=False)\n    parser.add_argument(\"--slerp\", type=bool, default=True)\n    parser.add_argument(\"--latents\", type=str, default=None)\n\n    args = parser.parse_args()\n\n    th.set_grad_enabled(False)\n    th.backends.cudnn.benchmark = True\n    mp.set_start_method(\"spawn\")\n\n    if args.stylegan1:\n        generator = G_style(output_size=args.out_size, checkpoint=args.ckpt).cuda()\n    else:\n        args.latent = 512\n        args.n_mlp = 8\n        generator = Generator(\n            args.G_res,\n            args.latent,\n            args.n_mlp,\n            channel_multiplier=args.channel_multiplier,\n            constant_input=args.const,\n            checkpoint=args.ckpt,\n            output_size=args.out_size,\n        ).cuda()\n    # generator = th.nn.DataParallel(generator)\n\n    if args.latents is not None:\n        styles = th.from_numpy(np.load(args.latents))\n    else:\n        # styles1 = th.randn((int(args.duration / 3), 512), device=\"cuda\")\n        # styles1 = generator(styles1, map_latents=True)\n        # styles2 = th.randn((int(args.duration / 3), 512), device=\"cuda\")\n        # styles2 = generator(styles2, map_latents=True)\n        # styles3 = th.randn((int(args.duration / 3), 512), device=\"cuda\")\n        # styles3 = generator(styles3, map_latents=True)\n\n        styles = th.randn((args.duration, 512), device=\"cuda\")\n        styles = generator(styles, map_latents=True)\n\n    latents = th.cat([styles[[0]]] * args.n_frames, axis=0)\n\n    # moving_low = spline_loops(\n    #     styles1.cpu(), 0, int(args.n_frames / 3), num_loops=1, smoothing=1, s=args.slerp\n    # ).cuda()[:, :5]\n    # moving_mid = spline_loops(\n    #     styles2.cpu(), 0, int(args.n_frames / 3), num_loops=1, smoothing=1, s=args.slerp\n    # ).cuda()[:, 5:10]\n    # moving_hi = spline_loops(\n    #     styles3.cpu(), 0, int(args.n_frames / 3), num_loops=1, smoothing=1, s=args.slerp\n    # ).cuda()[:, 10:]\n    # static_low = th.cat([moving_low[[0]]] * int(args.n_frames / 3), axis=0)\n    # static_mid = th.cat([moving_mid[[0]]] * int(args.n_frames / 3), axis=0)\n    # static_hi = th.cat([moving_hi[[0]]] * int(args.n_frames / 3), axis=0)\n\n    # print(\n    #     th.cat([moving_low, static_mid, static_hi], axis=1).shape,\n    #     th.cat([static_low[:60], static_mid[:60], static_hi[:60]], axis=1).shape,\n    #     th.cat([static_low, moving_mid, static_hi], axis=1).shape,\n    #     th.cat([static_low[:60], static_mid[:60], static_hi[:60]], axis=1).shape,\n    #     th.cat([static_low, static_mid, moving_hi], axis=1).shape,\n    # )\n\n    # print(th.cat([static_low[[0]], static_mid[[0]], static_hi[[0]]], axis=1).cpu().numpy().shape)\n    # np.save(\"latents_example.npy\", th.cat([static_low[[0]], static_mid[[0]], static_hi[[0]]], axis=1).cpu().numpy())\n\n    # latents = th.cat(\n    #     [\n    #         th.cat([moving_low, static_mid, static_hi], axis=1),\n    #         th.cat([static_low[:60], static_mid[:60], static_hi[:60]], axis=1),\n    #         th.cat([static_low, moving_mid, static_hi], axis=1),\n    #         th.cat([static_low[:60], static_mid[:60], static_hi[:60]], axis=1),\n    #         th.cat([static_low, static_mid, moving_hi], axis=1),\n    #     ],\n    #     axis=0,\n    # ).float()\n\n    # latents = gaussian_filter(latents, 7)\n\n    latents = latents.cpu()\n\n    print(\"latent shape: \")\n    print(latents.shape, \"\\n\")\n\n    log_max_res = int(np.log2(args.out_size))\n    log_min_res = 2 + (log_max_res - int(np.log2(args.G_res)))\n\n    noise = []\n    if args.stylegan1:\n        for s in range(log_min_res, log_max_res + 1):\n            h = 2 ** s\n            w = (2 if args.out_size == 1920 else 1) * 2 ** s\n            noise.append(th.randn((1, 1, h, w), device=\"cuda\"))\n    else:\n        for s in range(2 * log_min_res + 1, 2 * (log_max_res + 1), 1):\n            h = 2 ** int(s / 2)\n            w = (2 if args.out_size == 1920 else 1) * 2 ** int(s / 2)\n            noise.append(th.randn((1, 1, h, w), device=\"cuda\"))\n\n    def create_circular_mask(h, w, center=None, radius=None):\n        if center is None:  # use the middle of the image\n            center = (int(w / 2), int(h / 2))\n        if radius is None:  # use the smallest distance between the center and image walls\n            radius = min(center[0], center[1], w - center[0], h - center[1])\n        Y, X = np.ogrid[:h, :w]\n        dist_from_center = np.sqrt((X - center[0]) ** 2 + (Y - center[1]) ** 2)\n        mask = dist_from_center <= radius\n        return th.from_numpy(mask)\n\n    print(\"noise shapes: \")\n    for i, n in enumerate(noise):\n        if n is None:\n            continue\n        if i > 14:\n            noise[i] = None\n            continue\n\n        # mask = create_circular_mask(n.shape[-2], n.shape[-1], radius=n.shape[-1] / 2.5)[None, ...].float()\n        # mask = th.stack(\n        #     [\n        #         th.cat(\n        #             [\n        #                 th.zeros((int(n.shape[-2] * 1 / 2))),\n        #                 th.linspace(0, 1, int(n.shape[-2] * 1 / 4)),\n        #                 th.ones((int(n.shape[-2] * 1 / 4))),\n        #             ],\n        #             axis=0,\n        #         )\n        #     ]\n        #     * n.shape[-1]\n        # ).T[None, ...]\n        # mask = th.stack([mask] * n.shape[0], axis=0)\n        # noise[i] = mask * n[[0]].cpu()  # gaussian_filter(n, 24).cpu()\n\n        if i < 4:\n            moving = 2 * gaussian_filter(th.randn((200, 1, n.shape[-2], n.shape[-1]), device=\"cuda\"), 3)\n            # moving /= moving.std()\n            static = th.cat([n] * (len(latents) - len(moving)))\n            print(moving.shape, static.shape)\n            # static /= static.std()\n            noise[i] = th.cat([moving, static], axis=0)\n        elif 4 <= i < 8:\n            static1 = th.cat([n] * (260))\n            # static1 /= static1.std()\n            moving = 4 * gaussian_filter(th.randn((200, 1, n.shape[-2], n.shape[-1]), device=\"cuda\"), 3)\n            # moving /= moving.std()\n            static2 = th.cat([n] * (len(latents) - 460))\n            print(static1.shape, moving.shape, static2.shape)\n            # static2 /= static2.std()\n            noise[i] = th.cat([static1, moving, static2], axis=0)\n        elif i >= 8:\n            moving = 8 * gaussian_filter(th.randn((200, 1, n.shape[-2], n.shape[-1]), device=\"cuda\"), 3)\n            # moving /= moving.std()\n            static = th.cat([n] * (len(latents) - len(moving)))\n            print(static.shape, moving.shape)\n            # static /= static.std()\n            noise[i] = th.cat([static, moving], axis=0)\n        noise[i] = gaussian_filter(noise[i].cuda(), 7).cpu()\n\n        # noise[i] = th.cat([n[[0]]] * len(latents), axis=0).cpu()  # gaussian_filter(n, 24).cpu()\n        # noise[i] /= noise[i].std()\n\n        # if i > 2 and i < 13:\n        #     # xs = 8 * np.pi * th.linspace(0, 1, n.shape[-1])\n        #     # ys = th.linspace(0, 2 * np.pi, n.shape[-2])\n        #     # ts = 8 * np.pi * th.linspace(0, 1, n.shape[0])\n        #     # horiz = xs[None, None, None, :] + ys[None, None, :, None] + ts[:, None, None, None]\n        #     # vert = (\n        #     #     xs[None, None, None, :] / (4 * np.pi)\n        #     #     + 4 * np.pi * ys[None, None, :, None]\n        #     #     + 2 * ts[:, None, None, None]\n        #     # )\n        #     # moving_noise = th.sin(horiz.cuda() * vert.cuda() + n / 4)\n        #     # moving_noise = gaussian_filter(moving_noise, 6).cpu()\n        #     # moving_noise /= moving_noise.std() / 2\n        #     moving_noise = perlin_noise((n.shape[0], n.shape[-2], n.shape[-1]), (10, 8, 8))[:, None, ...]\n        #     moving_noise += gaussian_filter(n, 8) / 2.5\n        #     moving_noise /= moving_noise.std() / 1.5\n        #     noise[i] += (1 - mask) * moving_noise.cpu()\n\n        print(i, noise[i].shape, noise[i].std())\n    print()\n\n    import ffmpeg\n\n    output_name = f\"/home/hans/neurout/{args.ckpt.split('/')[-1].split('.')[0]}_{uuid.uuid4().hex[:8]}\"\n\n    video = (\n        ffmpeg.input(\"pipe:\", format=\"rawvideo\", pix_fmt=\"rgb24\", framerate=len(latents) / args.duration, s=\"256x256\")\n        .output(f\"{output_name}_noise.mp4\", framerate=len(latents) / args.duration, vcodec=\"libx264\", preset=\"slow\",)\n        .global_args(\"-benchmark\", \"-stats\", \"-hide_banner\")\n        .overwrite_output()\n        .run_async(pipe_stdin=True)\n    )\n    print(\n        noise[3][:200].shape,\n        noise[3][-45:-15].shape,\n        noise[7][15:45].shape,\n        noise[7][260:460].shape,\n        noise[7][15:45].shape,\n        noise[12][15:45].shape,\n        noise[12][520:].shape,\n    )\n    # output = noise[-5].permute(0, 2, 3, 1).numpy()\n    output = th.cat(\n        [\n            F.interpolate(noise[3][:200], (256, 256)),\n            F.interpolate(th.cat([noise[3][[200]]] * 30, axis=0), (256, 256)),\n            F.interpolate(th.cat([noise[7][[260]]] * 30, axis=0), (256, 256)),\n            F.interpolate(noise[7][260:460], (256, 256)),\n            F.interpolate(th.cat([noise[7][[460]]] * 30, axis=0), (256, 256)),\n            F.interpolate(th.cat([noise[12][[520]]] * 30, axis=0), (256, 256)),\n            F.interpolate(noise[12][520:], (256, 256)),\n        ],\n        axis=0,\n    )\n    print(output.shape)\n    output = output.permute(0, 2, 3, 1).numpy()\n    print(output.shape)\n    output = output / output.max()\n    output = output - output.min()\n    output = output * 255\n    output = output.astype(np.uint8)\n    output = np.concatenate([output] * 3, axis=3)\n    for frame in output:\n        video.stdin.write(frame.tobytes())\n    video.stdin.close()\n    video.wait()\n    # video.close()\n\n    # video = (\n    #     ffmpeg.input(\"pipe:\", format=\"rawvideo\", pix_fmt=\"rgb24\", framerate=len(latents) / args.duration, s=\"23x23\")\n    #     .output(f\"{output_name}_latents.mp4\", framerate=len(latents) / args.duration, vcodec=\"libx264\", preset=\"slow\",)\n    #     .global_args(\"-benchmark\", \"-stats\", \"-hide_banner\")\n    #     .overwrite_output()\n    #     .run_async(pipe_stdin=True)\n    # )\n    # output = th.cat(\n    #     [\n    #         latents[: int(args.n_frames / 3), 0],  #                                lo\n    #         latents[-45:-15, 0],  #                                                      pause lo\n    #         latents[15:45, 7],  #                                                       pause mid\n    #         latents[60 + int(args.n_frames / 3) : 60 + int(2 * args.n_frames / 3), 7],  #   mid\n    #         latents[15:45, 7],  #                                                       pause mid\n    #         latents[15:45, 14],  #                                                      pause hit\n    #         latents[120 + int(2 * args.n_frames / 3) :, 14],  #                           hi\n    #     ],\n    #     axis=0,\n    # )\n    # print(output.shape)\n    # output = th.cat([output, th.zeros((len(latents), 17))], axis=1)\n    # print(output.shape)\n    # output = output.reshape((len(latents), 23, 23, 1)).numpy()\n    # print(output.shape)\n    # output = output / output.max()\n    # output = output - output.min()\n    # output = output * 255\n    # output = output.astype(np.uint8)\n    # output = np.concatenate([output] * 3, axis=3)\n    # for frame in output:\n    #     video.stdin.write(frame.tobytes())\n    # video.stdin.close()\n    # video.wait()\n\n    # noise = []\n    # if args.stylegan1:\n    #     for s in range(log_min_res, log_max_res + 1):\n    #         h = 2 ** s\n    #         w = (2 if args.out_size == 1920 else 1) * 2 ** s\n    #         noise.append(np.random.normal(size=(args.n_frames, 1, h, w)))\n    # else:\n    #     for s in range(2 * log_min_res + 1, 2 * (log_max_res + 1), 1):\n    #         h = 2 ** int(s / 2)\n    #         w = (2 if args.out_size == 1920 else 1) * 2 ** int(s / 2)\n    #         noise.append(np.random.normal(size=(args.n_frames, 1, h, w)))\n\n    # print(\"noise shapes: \")\n    # for i, n in enumerate(noise):\n    #     if n is None:\n    #         continue\n    #     noise[i] = th.from_numpy(ndi.gaussian_filter(n, [15, 0, 0, 0], mode=\"wrap\"))\n    #     print(n.shape)\n    # print()\n\n    class addNoise(th.nn.Module):\n        def __init__(self, noise):\n            super(addNoise, self).__init__()\n            self.noise = noise\n\n        def forward(self, x):\n            return x + self.noise\n\n    manipulations = []\n    if log_min_res > 2:\n        reflects = []\n        for lres in range(2, log_min_res):\n            half = 2 ** (lres - 1)\n            reflects.append(th.nn.ReplicationPad2d((half, half, half, half)))\n        manipulations += [\n            {\n                \"layer\": 0,\n                \"transform\": th.nn.Sequential(\n                    *reflects, addNoise(2 * th.randn(size=(1, 1, 2 ** log_min_res, 2 ** log_min_res), device=\"cuda\")),\n                ),\n            }\n        ]\n\n    # tl = 4\n    # width = lambda s: (2 if args.out_size == 1920 else 1) * 2 ** int(s)\n    # translation = (\n    #     th.tensor([np.linspace(0, width(tl), args.n_frames + 1), np.zeros((args.n_frames + 1,))]).float().T[:-1]\n    # )\n    # manipulations += [{\"layer\": tl, \"transform\": \"translateX\", \"params\": translation}]\n\n    # zl = 6\n    # print(\n    #     th.cat(\n    #         [\n    #             th.linspace(-1, 3, int(args.n_frames / 2)),\n    #             th.linspace(3, -1, args.n_frames - int(args.n_frames / 2)) + 1,\n    #         ]\n    #     ).shape\n    # )\n    # zoom = gaussian_filter(\n    #     th.cat(\n    #         [\n    #             th.linspace(0, 3, int(args.n_frames / 2), dtype=th.float32, device=\"cuda\"),\n    #             th.linspace(3, 0, args.n_frames - int(args.n_frames / 2), dtype=th.float32, device=\"cuda\") + 1,\n    #         ]\n    #     )[:, None, None],\n    #     30,\n    # ).squeeze()\n    # zoom -= zoom.min()\n    # zoom /= zoom.max()\n    # # zoom *= 1.5\n    # zoom += 0.5\n    # print(zoom.min().item(), zoom.max().item(), zoom.shape)\n    # manipulations += [{\"layer\": zl, \"transform\": \"zoom\", \"params\": zoom}]\n\n    # rl = 6\n    # rotation = th.nn.Sigmoid()(th.tensor(np.linspace(0.0, 1.0, args.n_frames + 1), device=\"cuda\").float())\n    # rotation -= rotation.min()\n    # rotation /= rotation.max()\n    # rotation = rotation[:-1]\n    # manipulations += [{\"layer\": rl, \"transform\": \"rotate\", \"params\": (360.0 * rotation).cpu()}]\n\n    render(\n        generator=generator,\n        latents=latents,\n        noise=noise,\n        offset=0,\n        duration=args.duration,\n        batch_size=args.batch,\n        truncation=args.truncation,\n        manipulations=manipulations,\n        out_size=args.out_size,\n        output_file=f\"{output_name}.mp4\",\n    )\n"
  },
  {
    "path": "gpu_profile.py",
    "content": "import datetime\nimport linecache\nimport os\n\nos.environ[\"CUDA_LAUNCH_BLOCKING\"] = \"1\"\n\nfrom py3nvml import py3nvml\nimport torch\nimport socket\n\n# different settings\nprint_tensor_sizes = True\nuse_incremental = False\n\nif \"GPU_DEBUG\" in os.environ:\n    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\"\n    print(\"profiling gpu usage to \", gpu_profile_fn)\n\n## Global variables\nlast_tensor_sizes = set()\nlast_meminfo_used = 0\nlineno = None\nfunc_name = None\nfilename = None\nmodule_name = None\n\n\ndef gpu_profile(frame, event, arg):\n    # it is _about to_ execute (!)\n    global last_tensor_sizes\n    global last_meminfo_used\n    global lineno, func_name, filename, module_name\n\n    if event == \"line\":\n        try:\n            # about _previous_ line (!)\n            if lineno is not None:\n                py3nvml.nvmlInit()\n                handle = py3nvml.nvmlDeviceGetHandleByIndex(int(os.environ[\"GPU_DEBUG\"]))\n                meminfo = py3nvml.nvmlDeviceGetMemoryInfo(handle)\n                line = linecache.getline(filename, lineno)\n                where_str = module_name + \" \" + func_name + \":\" + str(lineno)\n\n                new_meminfo_used = meminfo.used\n                mem_display = new_meminfo_used - last_meminfo_used if use_incremental else new_meminfo_used\n                if abs(new_meminfo_used - last_meminfo_used) / 1024 ** 2 > 256:\n                    with open(gpu_profile_fn, \"a+\") as f:\n                        f.write(f\"{where_str:<50}\" f\":{(mem_display)/1024**2:<7.1f}Mb \" f\"{line.rstrip()}\\n\")\n\n                        last_meminfo_used = new_meminfo_used\n                        if print_tensor_sizes is True:\n                            for tensor in get_tensors():\n                                if not hasattr(tensor, \"dbg_alloc_where\"):\n                                    tensor.dbg_alloc_where = where_str\n                            new_tensor_sizes = {(type(x), tuple(x.size()), x.dbg_alloc_where) for x in get_tensors()}\n                            for t, s, loc in new_tensor_sizes - last_tensor_sizes:\n                                f.write(f\"+ {loc:<50} {str(s):<20} {str(t):<10}\\n\")\n                            for t, s, loc in last_tensor_sizes - new_tensor_sizes:\n                                f.write(f\"- {loc:<50} {str(s):<20} {str(t):<10}\\n\")\n                            last_tensor_sizes = new_tensor_sizes\n                py3nvml.nvmlShutdown()\n\n            # save details about line _to be_ executed\n            lineno = None\n\n            func_name = frame.f_code.co_name\n            filename = frame.f_globals[\"__file__\"]\n            if filename.endswith(\".pyc\") or filename.endswith(\".pyo\"):\n                filename = filename[:-1]\n            module_name = frame.f_globals[\"__name__\"]\n            lineno = frame.f_lineno\n\n            # only profile codes within the parent folder, otherwise there are too many function calls into other pytorch scripts\n            # need to modify the key words below to suit your case.\n            if \"maua-stylegan2\" not in os.path.dirname(os.path.abspath(filename)):\n                lineno = None  # skip current line evaluation\n\n            if (\n                \"car_datasets\" in filename\n                or \"_exec_config\" in func_name\n                or \"gpu_profile\" in module_name\n                or \"tee_stdout\" in module_name\n                or \"PIL\" in module_name\n            ):\n                lineno = None  # skip othe unnecessary lines\n\n            return gpu_profile\n\n        except (KeyError, AttributeError):\n            pass\n\n    return gpu_profile\n\n\ndef get_tensors(gpu_only=True):\n    import gc\n\n    for obj in gc.get_objects():\n        try:\n            if torch.is_tensor(obj):\n                tensor = obj\n            elif hasattr(obj, \"data\") and torch.is_tensor(obj.data):\n                tensor = obj.data\n            else:\n                continue\n\n            if tensor.is_cuda:\n                yield tensor\n        except Exception as e:\n            pass\n"
  },
  {
    "path": "gpumon.py",
    "content": "import argparse\nimport os\nimport signal\nimport subprocess\nimport time\nfrom queue import Empty, Queue\nfrom threading import Thread\n\nimport numpy as np\nimport wandb\n\nparser = argparse.ArgumentParser()\n\nparser.add_argument(\"--wbname\", type=str, required=True)\nparser.add_argument(\"--wbproj\", type=str, required=True)\nparser.add_argument(\"--wbgroup\", type=str, default=None)\n\nargs = parser.parse_args()\n\nif args.wbgroup is None:\n    wandb.init(project=args.wbproj, name=args.wbname, settings=wandb.Settings(_disable_stats=True))\nelse:\n    wandb.init(project=args.wbproj, group=args.wbgroup, name=args.wbname, settings=wandb.Settings(_disable_stats=True))\n\n\ndef enqueue_output(out, queue):\n    for line in iter(out.readline, b\"\"):\n        queue.put(line)\n    out.close()\n\n\nos.setpgrp()\n\nclock_proc = subprocess.Popen(\"nvidia-smi dmon -s c\", shell=True, stdout=subprocess.PIPE, bufsize=1)\nclock_proc.daemon = True\n\ntime.sleep(0.5)\n\nthrottle_reasons = [\n    \"clocks_throttle_reasons.gpu_idle\",\n    \"clocks_throttle_reasons.applications_clocks_setting\",\n    \"clocks_throttle_reasons.sw_power_cap\",\n    \"clocks_throttle_reasons.sw_thermal_slowdown\",\n    \"clocks_throttle_reasons.hw_slowdown\",\n    \"clocks_throttle_reasons.hw_thermal_slowdown\",\n    \"clocks_throttle_reasons.hw_power_brake_slowdown\",\n    \"clocks_throttle_reasons.sync_boost\",\n]\nthrottle_proc = subprocess.Popen(\n    f\"nvidia-smi --query-gpu=index,{','.join(throttle_reasons)} --format=csv,noheader --loop=1\",\n    shell=True,\n    stdout=subprocess.PIPE,\n    bufsize=1,\n)\nthrottle_proc.daemon = True\n\n# create queue that gets the output lines from both processes\nq = Queue()\nclock_thread = Thread(target=enqueue_output, args=(clock_proc.stdout, q))\nclock_thread.daemon = True\nthottle_thread = Thread(target=enqueue_output, args=(throttle_proc.stdout, q))\nthottle_thread.daemon = True\n\nclock_thread.start()\nthottle_thread.start()\n\nthrottles = [[], []]\nclocks = [[], []]\nwhile clock_proc.poll() is None or not q.empty():\n    try:\n        line = q.get_nowait()\n    except Empty:\n        pass\n    else:\n        line = line.decode(\"utf-8\").strip()\n        if \"#\" in line:\n            continue\n        if \",\" in line:\n            raw = line.split(\",\")\n            gpu = int(raw[0])\n            bits = [0 if \"Not\" in a else 1 for a in raw[1:]]\n            throttles[gpu].append(bits)\n            # print(gpu, bits)\n        else:\n            raw = line.split(\"  \")\n            gpu = int(raw[0])\n            clock = int(raw[-1])\n            clocks[gpu].append(clock)\n            # print(gpu, clock)\n\n    if len(clocks[0]) > 30:\n        try:\n            throttles = np.array(throttles)\n            clocks = np.array(clocks)\n            log_dict = {}\n            for gpu in [0, 1]:\n                log_dict[f\"gpu.{gpu}.clock.speed\"] = np.mean(clocks[gpu])\n\n                for r, reason in enumerate(throttle_reasons):\n                    log_dict[f\"gpu.{gpu}.{reason}\"] = np.mean(throttles[gpu, :, r])\n\n            print(\"\\n\".join([k.ljust(80) + str(v) for k, v in log_dict.items()]))\n            wandb.log(log_dict)\n        except:\n            pass\n\n        throttles = [[], []]\n        clocks = [[], []]\n\nos.kill(throttle_proc.pid, signal.SIGINT)\nos.kill(clock_proc.pid, signal.SIGINT)\n"
  },
  {
    "path": "lightning.py",
    "content": "import os\nimport gc\nimport math\nimport wandb\nimport random\nimport argparse\nimport validation\nimport torch as th\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torchvision as tv\nimport torchvision.transforms as transforms\nfrom collections import OrderedDict\nfrom torch.utils import data\nfrom dataset import MultiResolutionDataset\nfrom model import Generator, Discriminator\nimport pytorch_lightning as pl\n\n\ndef requires_grad(model, flag=True):\n    for p in model.parameters():\n        p.requires_grad = flag\n\n\ndef get_spectral_norms(model):\n    spectral_norms = {}\n    for name, param in model.named_parameters():\n        if param.numel() > 0:\n            spectral_norms[name] = nn.utils.spectral_norm(param)\n    return spectral_norms\n\n\nclass StyleGAN2(pl.LightningModule):\n    def __init__(self, hparams):\n        super().__init__()\n        self.hparams = hparams  # for automatic param saving with lightning\n        [setattr(self, k, v) for k, v in vars(hparams).items()]  # for easy access within module\n\n        self.generator = Generator(self.size, self.latent_size, self.n_mlp, channel_multiplier=self.channel_multiplier)\n        self.g_ema = Generator(self.size, self.latent_size, self.n_mlp, channel_multiplier=self.channel_multiplier)\n        self.g_ema.eval()\n        self.accumulate_g(0)\n\n        self.discriminator = Discriminator(self.size, channel_multiplier=self.channel_multiplier)\n\n        self.sample_z = th.randn(self.n_sample, self.latent_size)\n\n        self.mean_path_length = th.tensor(0.0)\n\n    def forward(self, z):\n        return self.generator(z)\n\n    def accumulate_g(self, decay=0.5 ** (32.0 / (10_000))):\n        par1 = dict(self.g_ema.named_parameters())\n        par2 = dict(self.generator.named_parameters())\n        for name, param in self.g_ema.named_parameters():\n            param.data = decay * par1[name].data + (1 - decay) * par2[name].data\n\n    def configure_optimizers(self):\n        g_reg_ratio = self.g_reg_every / (self.g_reg_every + 1)\n        d_reg_ratio = self.d_reg_every / (self.d_reg_every + 1)\n\n        g_optim = th.optim.Adam(\n            self.generator.parameters(), lr=self.lr * g_reg_ratio, betas=(0 ** g_reg_ratio, 0.99 ** g_reg_ratio),\n        )\n        d_optim = th.optim.Adam(\n            self.discriminator.parameters(), lr=self.lr * d_reg_ratio, betas=(0 ** d_reg_ratio, 0.99 ** d_reg_ratio),\n        )\n\n        return [g_optim, g_optim, d_optim, d_optim], []\n\n    def configure_apex(self, amp, model, optimizers, amp_level):\n        amp_optimizers = []\n        for optimizer in optimizers:\n            try:\n                amp_model, amp_optimizer = amp.initialize(model, optimizer, opt_level=amp_level,)\n            except RuntimeError as err:\n                print(err)\n                print(\"Skipping this optimizer\")\n            amp_optimizers.append(amp_optimizer)\n        return amp_model, amp_optimizers\n\n    def train_dataloader(self):\n        transform = transforms.Compose(\n            [\n                transforms.RandomVerticalFlip(p=0.5 if self.vflip else 0),\n                transforms.RandomHorizontalFlip(p=0.5 if self.hflip else 0),\n                transforms.ToTensor(),\n                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True),\n            ]\n        )\n        dataset = MultiResolutionDataset(self.path, transform, self.size)\n        loader = data.DataLoader(dataset, batch_size=self.batch_size, shuffle=True, num_workers=0)\n        return loader\n\n    def d_logistic_loss(self, real_pred, fake_pred):\n        real_loss = F.softplus(-real_pred)\n        fake_loss = F.softplus(fake_pred)\n        return real_loss.mean() + fake_loss.mean()\n\n    def d_r1_loss(self, real_pred, real_img):\n        (grad_real,) = th.autograd.grad(outputs=real_pred.sum(), inputs=real_img, create_graph=True)\n        grad_penalty = grad_real.pow(2).view(grad_real.shape[0], -1).sum(1).mean()\n        return grad_penalty\n\n    def g_nonsaturating_loss(self, fake_pred):\n        loss = F.softplus(-fake_pred).mean()\n        return loss\n\n    def g_path_regularize(self, fake_img, latents, mean_path_length, decay=0.01):\n        noise = th.randn_like(fake_img) / math.sqrt(fake_img.shape[2] * fake_img.shape[3])\n        # print(fake_img.requires_grad, noise.requires_grad, latents.requires_grad)\n        (grad,) = th.autograd.grad(outputs=(fake_img * noise).sum(), inputs=latents, create_graph=True)\n        path_lengths = th.sqrt(grad.pow(2).sum(2).mean(1))\n        path_mean = mean_path_length + decay * (path_lengths.mean() - mean_path_length)\n        path_penalty = (path_lengths - path_mean).pow(2).mean()\n        return path_penalty, path_mean.detach(), path_lengths\n\n    def make_noise(self, batch, batch_size=None):\n        if batch_size is None:\n            batch_size = batch.size(0)\n        if self.mixing_prob > 0 and random.random() < self.mixing_prob:\n            return th.randn(2, batch_size, self.latent_size).type_as(batch).unbind(0)\n        else:\n            return [th.randn(batch_size, self.latent_size).type_as(batch)]\n\n    def training_step(self, real_img, batch_idx, optimizer_idx):\n        # real_img = real_img.half()\n        log_dict = {}\n\n        # train generator\n        if optimizer_idx == 0:\n            requires_grad(self.generator, True)\n            requires_grad(self.discriminator, False)\n\n            noise = self.make_noise(real_img)\n            # print(real_img.dtype, noise[0].dtype, real_img.device)\n            fake_img, _ = self.generator(noise)\n            fake_pred = self.discriminator(fake_img)\n            g_loss = self.g_nonsaturating_loss(fake_pred)\n\n            log_dict[\"Generator\"] = g_loss\n            # log_dict[\"Spectral Norms/Generator\"] = get_spectral_norms(self.generator)\n\n            # print(g_loss)\n\n            return OrderedDict({\"loss\": g_loss, \"log\": log_dict})\n\n        # maybe regularize generator\n        if optimizer_idx == 1:\n            if batch_idx % self.g_reg_every == 0:\n                path_batch_size = max(1, self.batch_size // self.path_batch_shrink)\n                noise = self.make_noise(real_img, path_batch_size)\n                fake_img, latents = self.generator(noise, return_latents=True)\n                path_loss, self.mean_path_length, path_lengths = self.g_path_regularize(\n                    fake_img, latents, self.mean_path_length.type_as(real_img)\n                )\n                weighted_path_loss = self.path_regularize * self.g_reg_every * path_loss\n                if self.path_batch_shrink:\n                    weighted_path_loss += 0 * fake_img[0, 0, 0, 0]\n\n                log_dict[\"Path Length Regularization\"] = path_loss\n                log_dict[\"Mean Path Length\"] = path_lengths.mean()\n\n                return OrderedDict({\"loss\": weighted_path_loss, \"log\": log_dict})\n            return OrderedDict({\"loss\": th.tensor(-69).type_as(real_img)})\n\n        # train discriminator\n        if optimizer_idx == 2:\n            requires_grad(self.generator, False)\n            requires_grad(self.discriminator, True)\n\n            noise = self.make_noise(real_img)\n            fake_img, _ = self.generator(noise)\n            fake_pred = self.discriminator(fake_img)\n\n            real_pred = self.discriminator(real_img)\n            d_loss = self.d_logistic_loss(real_pred, fake_pred)\n\n            log_dict[\"Discriminator\"] = d_loss\n            log_dict[\"Real Score\"] = real_pred.mean()\n            log_dict[\"Fake Score\"] = fake_pred.mean()\n            # log_dict[\"Spectral Norms/Discriminator\"] = get_spectral_norms(self.discriminator)\n\n            # print(d_loss)\n\n            return OrderedDict({\"loss\": d_loss, \"log\": log_dict})\n\n        # maybe regularize discriminator\n        if optimizer_idx == 3:\n            if batch_idx % self.d_reg_every == 0:\n                real_img.requires_grad = True\n                real_pred = self.discriminator(real_img)\n                r1_loss = self.d_r1_loss(real_pred, real_img)\n                weighted_r1_loss = self.r1 / 2 * r1_loss * self.d_reg_every + 0 * real_pred[0]\n\n                log_dict[\"R1\"] = r1_loss\n\n                return OrderedDict({\"loss\": weighted_r1_loss, \"log\": log_dict})\n            return OrderedDict({\"loss\": th.tensor(-69).type_as(real_img)})\n\n    def backward(self, trainer, loss, optimizer, optimizer_idx):\n        if optimizer_idx == 0:\n            super(StyleGAN2, self).backward(trainer, loss, optimizer, optimizer_idx)\n        if optimizer_idx == 1 and loss != -69:\n            super(StyleGAN2, self).backward(trainer, loss, optimizer, optimizer_idx)\n        if optimizer_idx == 2:\n            super(StyleGAN2, self).backward(trainer, loss, optimizer, optimizer_idx)\n        if optimizer_idx == 3 and loss != -69:\n            super(StyleGAN2, self).backward(trainer, loss, optimizer, optimizer_idx)\n\n    def optimizer_step(self, cur_epoch, batch_idx, optimizer, optimizer_idx, closure):\n        if optimizer_idx == 0:\n            super(StyleGAN2, self).optimizer_step(cur_epoch, batch_idx, optimizer, optimizer_idx, closure)\n        if optimizer_idx == 1:\n            if batch_idx % self.g_reg_every == 0:\n                super(StyleGAN2, self).optimizer_step(cur_epoch, batch_idx, optimizer, optimizer_idx, closure)\n            self.accumulate_g()\n        if optimizer_idx == 2:\n            super(StyleGAN2, self).optimizer_step(cur_epoch, batch_idx, optimizer, optimizer_idx, closure)\n        if optimizer_idx == 3 and batch_idx % self.d_reg_every == 0:\n            super(StyleGAN2, self).optimizer_step(cur_epoch, batch_idx, optimizer, optimizer_idx, closure)\n\n    def prepare_data(self):\n        validation.get_dataset_inception_features(self.train_dataloader(), self.name, self.size)\n\n    def val_dataloader(self):\n        return [[th.arange(0, 1)]]\n\n    def validation_step(self, batch, batch_idx):\n        # gc.collect()\n        # th.cuda.empty_cache()\n        # output = OrderedDict({\"FID\": th.tensor(-69).type_as(batch), \"PPL\": th.tensor(-69).type_as(batch)})\n        # for task in batch:\n        #     # if task == 1:\n        #     output[\"FID\"] = fid.validation_fid(\n        #         self.g_ema.to(batch.device), self.val_batch_size, self.fid_n_sample, self.fid_truncation, self.name,\n        #     )\n        #     # if task == 0:\n        #     output[\"PPL\"] = ppl.validation_ppl(\n        #         self.g_ema.to(batch.device),\n        #         self.val_batch_size,\n        #         self.ppl_n_sample,\n        #         self.ppl_space,\n        #         self.ppl_crop,\n        #         self.latent_size,\n        #     )\n        return OrderedDict({\"batch\": batch})  # output\n\n    def validation_epoch_end(self, outputs):\n        batch = outputs[0][\"batch\"]\n        gc.collect()\n        th.cuda.empty_cache()\n        val_fid = validation.fid(\n            self.g_ema.to(batch.device), self.val_batch_size, self.fid_n_sample, self.fid_truncation, self.name,\n        )[\"FID\"]\n        val_ppl = validation.ppl(\n            self.g_ema.to(batch.device),\n            self.val_batch_size,\n            self.ppl_n_sample,\n            self.ppl_space,\n            self.ppl_crop,\n            self.latent_size,\n        )\n        with th.no_grad():\n            self.g_ema.eval()\n            sample, _ = self.g_ema([self.sample_z.to(next(self.g_ema.parameters()).device)])\n            grid = tv.utils.make_grid(\n                sample, nrow=int(round(4.0 / 3 * self.n_sample ** 0.5)), normalize=True, range=(-1, 1)\n            )\n            self.logger.experiment.log(\n                {\"Generated Images EMA\": [wandb.Image(grid, caption=f\"Step {self.global_step}\")]}\n            )\n\n            self.generator.eval()\n            sample, _ = self.generator([self.sample_z.to(next(self.generator.parameters()).device)])\n            grid = tv.utils.make_grid(\n                sample, nrow=int(round(4.0 / 3 * self.n_sample ** 0.5)), normalize=True, range=(-1, 1)\n            )\n            self.logger.experiment.log({\"Generated Images\": [wandb.Image(grid, caption=f\"Step {self.global_step}\")]})\n            self.generator.train()\n\n        # val_fid = [score for score in outputs[0][\"FID\"] if score != -69][0]\n        # val_ppl = [score for score in outputs[0][\"PPL\"] if score != -69][0]\n\n        gc.collect()\n        th.cuda.empty_cache()\n\n        return {\"val_loss\": val_fid, \"log\": {\"Validation/FID\": val_fid, \"Validation/PPL\": val_ppl}}\n\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser()\n\n    # data options\n    parser.add_argument(\"path\", type=str)\n    parser.add_argument(\"--vflip\", type=bool, default=False)\n    parser.add_argument(\"--hflip\", type=bool, default=True)\n\n    # training options\n    parser.add_argument(\"--epochs\", type=int, default=100)\n    parser.add_argument(\"--batch_size\", type=int, default=16)\n    parser.add_argument(\"--checkpoint\", type=str, default=None)\n\n    # model options\n    parser.add_argument(\"--latent_size\", type=int, default=512)\n    parser.add_argument(\"--n_mlp\", type=int, default=8)\n    parser.add_argument(\"--n_sample\", type=int, default=32)\n    parser.add_argument(\"--size\", type=int, default=256)\n\n    # loss options\n    parser.add_argument(\"--r1\", type=float, default=10)\n    parser.add_argument(\"--path_regularize\", type=float, default=2)\n    parser.add_argument(\"--path_batch_shrink\", type=int, default=2)\n    parser.add_argument(\"--d_reg_every\", type=int, default=16)\n    parser.add_argument(\"--g_reg_every\", type=int, default=4)\n    parser.add_argument(\"--mixing_prob\", type=float, default=0.9)\n    parser.add_argument(\"--lr\", type=float, default=0.002)\n    parser.add_argument(\"--channel_multiplier\", type=int, default=2)\n\n    # validation / logging options\n    parser.add_argument(\"--wandb\", type=bool, default=True)\n    parser.add_argument(\"--validation_interval\", type=float, default=0.25)\n    parser.add_argument(\"--val_batch_size\", type=int, default=24)\n    parser.add_argument(\"--fid_n_sample\", type=int, default=10000)\n    parser.add_argument(\"--fid_truncation\", type=float, default=0.7)\n    parser.add_argument(\"--ppl_space\", choices=[\"z\", \"w\"], default=\"w\")\n    parser.add_argument(\"--ppl_n_sample\", type=int, default=5000)\n    parser.add_argument(\"--ppl_crop\", type=bool, default=False)\n\n    # DevOps options\n    parser.add_argument(\"--num_gpus\", type=int, default=2)\n    parser.add_argument(\"--cudnn_benchmark\", type=bool, default=True)\n    parser.add_argument(\"--distributed_backend\", type=str, default=\"ddp\")\n\n    args = parser.parse_args()\n\n    args.name = os.path.splitext(os.path.basename(args.path))[0]\n\n    stylegan2 = StyleGAN2(args)\n    stylegan2.prepare_data()\n    stylegan2.train_dataloader()\n\n    checkpoint_callback = pl.callbacks.ModelCheckpoint(\n        filepath=\"/home/hans/modelzoo/maua-sg2/\" + args.name + \"-{epoch}-{val_loss:.0f}\", save_top_k=10\n    )\n\n    wandb_logger = pl.loggers.WandbLogger(project=\"maua-stylegan\")\n    # print(wandb_logger.experiment)\n\n    trainer = pl.Trainer(\n        gpus=args.num_gpus,\n        max_epochs=args.epochs,\n        logger=wandb_logger,\n        checkpoint_callback=checkpoint_callback,\n        early_stop_callback=None,\n        distributed_backend=args.distributed_backend,\n        benchmark=args.cudnn_benchmark,\n        val_check_interval=args.validation_interval,\n        num_sanity_val_steps=0,\n        terminate_on_nan=True,\n        resume_from_checkpoint=args.checkpoint,\n        amp_level=\"O2\",\n        precision=16,\n    )\n    trainer.fit(stylegan2)\n"
  },
  {
    "path": "lookahead_minimax.py",
    "content": "from collections import defaultdict\n\nimport torch\nfrom torch.optim.optimizer import Optimizer\n\n\nclass LookaheadMinimax(Optimizer):\n    r\"\"\"\n    A PyTorch implementation of the lookahead wrapper for GANs.\n\n    This optimizer performs the lookahead step on both the discriminator and generator optimizers after the generator's\n    optimizer takes a step. This ensures that joint minimax lookahead is used rather than alternating minimax lookahead\n    (which would result from simply applying the original Lookahead Optimizer to both networks separately).\n\n    Lookahead Minimax Optimizer: https://arxiv.org/abs/2006.14567\n    Lookahead Optimizer: https://arxiv.org/abs/1907.08610\n    \"\"\"\n\n    def __init__(self, G_optimizer, D_optimizer, la_steps=5, la_alpha=0.5, pullback_momentum=\"none\", accumulate=1):\n        \"\"\"\n        G_optimizer: generator optimizer\n        D_optimizer: discriminator optimizer\n        la_steps (int): number of lookahead steps\n        la_alpha (float): linear interpolation factor. 1.0 recovers the inner optimizer.\n        pullback_momentum (str): change to inner optimizer momentum on interpolation update\n        acumulate (int): number of gradient accumulation steps\n        \"\"\"\n        self.G_optimizer = G_optimizer\n        self.D_optimizer = D_optimizer\n\n        self._la_step = 0  # counter for inner optimizer\n        self.la_alpha = la_alpha\n        self._total_la_steps = la_steps * accumulate\n        self._la_steps = la_steps\n\n        pullback_momentum = pullback_momentum.lower()\n        assert pullback_momentum in [\"reset\", \"pullback\", \"none\"]\n        self.pullback_momentum = pullback_momentum\n\n        self.state = defaultdict(dict)\n\n        # Cache the current optimizer parameters\n        for group in G_optimizer.param_groups:\n            for p in group[\"params\"]:\n                param_state = self.state[p]\n                param_state[\"cached_G_params\"] = torch.zeros_like(p.data)\n                param_state[\"cached_G_params\"].copy_(p.data)\n                if self.pullback_momentum == \"pullback\":\n                    param_state[\"cached_G_mom\"] = torch.zeros_like(p.data)\n\n        for group in D_optimizer.param_groups:\n            for p in group[\"params\"]:\n                param_state = self.state[p]\n                param_state[\"cached_D_params\"] = torch.zeros_like(p.data)\n                param_state[\"cached_D_params\"].copy_(p.data)\n                if self.pullback_momentum == \"pullback\":\n                    param_state[\"cached_D_mom\"] = torch.zeros_like(p.data)\n\n    def __getstate__(self):\n        return {\n            \"state\": self.state,\n            \"G_optimizer\": self.G_optimizer,\n            \"D_optimizer\": self.D_optimizer,\n            \"la_alpha\": self.la_alpha,\n            \"_la_step\": self._la_step,\n            \"_total_la_steps\": self._la_steps,\n            \"pullback_momentum\": self.pullback_momentum,\n        }\n\n    def zero_grad(self):\n        self.G_optimizer.zero_grad()\n\n    def get_la_step(self):\n        return self._la_step\n\n    def state_dict(self):\n        return self.G_optimizer.state_dict()\n\n    def load_state_dict(self, G_state_dict, D_state_dict):\n        self.G_optimizer.load_state_dict(G_state_dict)\n        self.D_optimizer.load_state_dict(D_state_dict)\n\n        # Cache the current optimizer parameters\n        for group in self.G_optimizer.param_groups:\n            for p in group[\"params\"]:\n                param_state = self.state[p]\n                param_state[\"cached_G_params\"] = torch.zeros_like(p.data)\n                param_state[\"cached_G_params\"].copy_(p.data)\n                if self.pullback_momentum == \"pullback\":\n                    param_state[\"cached_G_mom\"] = self.G_optimizer.state[p][\"momentum_buffer\"]\n\n        for group in self.D_optimizer.param_groups:\n            for p in group[\"params\"]:\n                param_state = self.state[p]\n                param_state[\"cached_D_params\"] = torch.zeros_like(p.data)\n                param_state[\"cached_D_params\"].copy_(p.data)\n                if self.pullback_momentum == \"pullback\":\n                    param_state[\"cached_D_mom\"] = self.D_optimizer.state[p][\"momentum_buffer\"]\n\n    def _backup_and_load_cache(self):\n        \"\"\"\n        Useful for performing evaluation on the slow weights (which typically generalize better)\n        \"\"\"\n        for group in self.G_optimizer.param_groups:\n            for p in group[\"params\"]:\n                param_state = self.state[p]\n                param_state[\"backup_G_params\"] = torch.zeros_like(p.data)\n                param_state[\"backup_G_params\"].copy_(p.data)\n                p.data.copy_(param_state[\"cached_G_params\"])\n\n        for group in self.D_optimizer.param_groups:\n            for p in group[\"params\"]:\n                param_state = self.state[p]\n                param_state[\"backup_D_params\"] = torch.zeros_like(p.data)\n                param_state[\"backup_D_params\"].copy_(p.data)\n                p.data.copy_(param_state[\"cached_D_params\"])\n\n    def _clear_and_load_backup(self):\n        for group in self.G_optimizer.param_groups:\n            for p in group[\"params\"]:\n                param_state = self.state[p]\n                p.data.copy_(param_state[\"backup_G_params\"])\n                del param_state[\"backup_G_params\"]\n\n        for group in self.D_optimizer.param_groups:\n            for p in group[\"params\"]:\n                param_state = self.state[p]\n                p.data.copy_(param_state[\"backup_D_params\"])\n                del param_state[\"backup_D_params\"]\n\n    @property\n    def param_groups(self):\n        return self.G_optimizer.param_groups\n\n    def step(self, closure=None):\n        \"\"\"\n        Performs a single Lookahead optimization step on BOTH optimizers after the generator's optimizer step.\n\n        This allows the discriminator's optimizer to take more steps when using a higher step ratio and still have the\n        lookahead step being performed once after k generator steps. This also ensures the optimizers are updated with\n        the lookahead step simultaneously, rather than in alternating fashion.\n\n        Arguments:\n            closure (callable, optional): A closure that reevaluates the model\n                and returns the loss.\n        \"\"\"\n        loss = self.G_optimizer.step(closure)\n        self._la_step += 1\n\n        if self._la_step >= self._total_la_steps:\n            with torch.cuda.amp.autocast(enabled=False):\n                self._la_step = 0\n\n                # Lookahead and cache the current generator optimizer parameters\n                for group in self.G_optimizer.param_groups:\n                    for p in group[\"params\"]:\n                        param_state = self.state[p]\n                        p.data.mul_(self.la_alpha).add_(1.0 - self.la_alpha, param_state[\"cached_G_params\"])\n                        param_state[\"cached_G_params\"].copy_(p.data)\n\n                        if self.pullback_momentum == \"pullback\":\n                            internal_momentum = self.G_optimizer.state[p][\"momentum_buffer\"]\n                            self.G_optimizer.state[p][\"momentum_buffer\"] = internal_momentum.mul_(self.la_alpha).add_(\n                                1.0 - self.la_alpha, param_state[\"cached_G_mom\"]\n                            )\n                            param_state[\"cached_G_mom\"] = self.G_optimizer.state[p][\"momentum_buffer\"]\n                        elif self.pullback_momentum == \"reset\":\n                            self.G_optimizer.state[p][\"momentum_buffer\"] = torch.zeros_like(p.data)\n\n                # Lookahead and cache the current discriminator optimizer parameters\n                for group in self.D_optimizer.param_groups:\n                    for p in group[\"params\"]:\n                        param_state = self.state[p]\n                        p.data.mul_(self.la_alpha).add_(1.0 - self.la_alpha, param_state[\"cached_D_params\"])\n                        param_state[\"cached_D_params\"].copy_(p.data)\n\n                        if self.pullback_momentum == \"pullback\":\n                            internal_momentum = self.D_optimizer.state[p][\"momentum_buffer\"]\n                            self.D_optimizer.state[p][\"momentum_buffer\"] = internal_momentum.mul_(self.la_alpha).add_(\n                                1.0 - self.la_alpha, param_state[\"cached_D_mom\"]\n                            )\n                            param_state[\"cached_D_mom\"] = self.optimizer.state[p][\"momentum_buffer\"]\n                        elif self.pullback_momentum == \"reset\":\n                            self.D_optimizer.state[p][\"momentum_buffer\"] = torch.zeros_like(p.data)\n\n        return loss\n"
  },
  {
    "path": "lucidrains.py",
    "content": "import json, time, pickle, argparse\nfrom math import floor, log2, sqrt\nfrom random import random\nfrom shutil import rmtree\nfrom functools import partial\nfrom datetime import datetime\nimport multiprocessing\nfrom PIL import Image\nfrom pathlib import Path\nfrom retry.api import retry_call\nfrom tqdm import tqdm\n\nimport numpy as np\nfrom scipy import linalg\n\nimport torch\nfrom torch import nn\nfrom torch.utils import data\nimport torch.nn.functional as F\n\nfrom torch_optimizer import DiffGrad\nfrom torch.autograd import grad as torch_grad\n\nimport torchvision\nfrom torchvision import transforms\n\nfrom vector_quantize_pytorch import VectorQuantize\nfrom linear_attention_transformer import ImageLinearAttention\nfrom contrastive_learner import ContrastiveLearner, RandomApply\n\nfrom kornia import augmentation as augs\nfrom kornia import filters\n\nimport validation\nfrom validation.inception import InceptionV3\nfrom validation import lpips\n\nimport wandb\n\ntry:\n    from apex import amp\n\n    APEX_AVAILABLE = True\nexcept:\n    APEX_AVAILABLE = False\n\nassert torch.cuda.is_available(), \"You need to have an Nvidia GPU with CUDA installed.\"\n\nnum_cores = multiprocessing.cpu_count()\n\n# constants\n\nEXTS = [\"jpg\", \"png\"]\nEPS = 1e-8\n\n# helper classes\n\n\nclass NanException(Exception):\n    pass\n\n\nclass EMA:\n    def __init__(self, beta):\n        super().__init__()\n        self.beta = beta\n\n    def update_average(self, old, new):\n        if old is None:\n            return new\n        return old * self.beta + (1 - self.beta) * new\n\n\nclass Flatten(nn.Module):\n    def forward(self, x):\n        return x.reshape(x.shape[0], -1)\n\n\nclass Residual(nn.Module):\n    def __init__(self, fn):\n        super().__init__()\n        self.fn = fn\n\n    def forward(self, x):\n        return self.fn(x) + x\n\n\nclass Rezero(nn.Module):\n    def __init__(self, fn):\n        super().__init__()\n        self.fn = fn\n        self.g = nn.Parameter(torch.zeros(1))\n\n    def forward(self, x):\n        return self.fn(x) * self.g\n\n\nclass PermuteToFrom(nn.Module):\n    def __init__(self, fn):\n        super().__init__()\n        self.fn = fn\n\n    def forward(self, x):\n        x = x.permute(0, 2, 3, 1)\n        out, loss = self.fn(x)\n        out = out.permute(0, 3, 1, 2)\n        return out, loss\n\n\n# helpers\n\n\ndef default(value, d):\n    return d if value is None else value\n\n\ndef cycle(iterable):\n    while True:\n        for i in iterable:\n            yield i\n\n\ndef cast_list(el):\n    return el if isinstance(el, list) else [el]\n\n\ndef is_empty(t):\n    if isinstance(t, torch.Tensor):\n        return t.nelement() == 0\n    return t is None\n\n\ndef raise_if_nan(t):\n    if torch.isnan(t):\n        raise NanException\n\n\ndef loss_backwards(fp16, loss, optimizer, **kwargs):\n    if fp16:\n        with amp.scale_loss(loss, optimizer) as scaled_loss:\n            scaled_loss.backward(**kwargs)\n    else:\n        loss.backward(**kwargs)\n\n\ndef gradient_penalty(images, output, weight=10):\n    batch_size = images.shape[0]\n    gradients = torch_grad(\n        outputs=output,\n        inputs=images,\n        grad_outputs=torch.ones(output.size()).cuda(),\n        create_graph=True,\n        retain_graph=True,\n        only_inputs=True,\n    )[0]\n\n    gradients = gradients.view(batch_size, -1)\n    return weight * ((gradients.norm(2, dim=1) - 1) ** 2).mean()\n\n\ndef noise(n, latent_dim):\n    return torch.randn(n, latent_dim).cuda()\n\n\ndef noise_list(n, layers, latent_dim):\n    return [(noise(n, latent_dim), layers)]\n\n\ndef mixed_list(n, layers, latent_dim):\n    tt = int(torch.rand(()).numpy() * layers)\n    return noise_list(n, tt, latent_dim) + noise_list(n, layers - tt, latent_dim)\n\n\ndef latent_to_w(style_vectorizer, latent_descr):\n    return [(style_vectorizer(z), num_layers) for z, num_layers in latent_descr]\n\n\ndef image_noise(n, im_size):\n    return torch.FloatTensor(n, im_size, im_size, 1).uniform_(0.0, 1.0).cuda()\n\n\ndef leaky_relu(p=0.2):\n    return nn.LeakyReLU(p, inplace=True)\n\n\ndef evaluate_in_chunks(max_batch_size, model, *args):\n    split_args = list(zip(*list(map(lambda x: x.split(max_batch_size, dim=0), args))))\n    chunked_outputs = [model(*i) for i in split_args]\n    if len(chunked_outputs) == 1:\n        return chunked_outputs[0]\n    return torch.cat(chunked_outputs, dim=0)\n\n\ndef styles_def_to_tensor(styles_def):\n    return torch.cat([t[:, None, :].expand(-1, n, -1) for t, n in styles_def], dim=1)\n\n\ndef set_requires_grad(model, bool):\n    for p in model.parameters():\n        p.requires_grad = bool\n\n\n# dataset\n\n\ndef convert_rgb_to_transparent(image):\n    if image.mode == \"RGB\":\n        return image.convert(\"RGBA\")\n    return image\n\n\ndef convert_transparent_to_rgb(image):\n    if image.mode == \"RGBA\":\n        return image.convert(\"RGB\")\n    return image\n\n\nclass expand_greyscale(object):\n    def __init__(self, num_channels):\n        self.num_channels = num_channels\n\n    def __call__(self, tensor):\n        return tensor.expand(self.num_channels, -1, -1)\n\n\ndef resize_to_minimum_size(min_size, image):\n    if max(*image.size) < min_size:\n        return torchvision.transforms.functional.resize(image, min_size)\n    return image\n\n\nclass Dataset(data.Dataset):\n    def __init__(self, folder, image_size, transparent=False):\n        super().__init__()\n        self.folder = folder\n        self.image_size = image_size\n        self.paths = [p for ext in EXTS for p in Path(f\"{folder}\").glob(f\"**/*.{ext}\")]\n\n        convert_image_fn = convert_transparent_to_rgb if not transparent else convert_rgb_to_transparent\n        num_channels = 3 if not transparent else 4\n\n        self.transform = transforms.Compose(\n            [\n                transforms.Lambda(convert_image_fn),\n                transforms.Lambda(partial(resize_to_minimum_size, image_size)),\n                transforms.RandomHorizontalFlip(),\n                transforms.Resize(image_size),\n                transforms.CenterCrop(image_size),\n                transforms.ToTensor(),\n                transforms.Lambda(expand_greyscale(num_channels)),\n            ]\n        )\n\n    def __len__(self):\n        return len(self.paths)\n\n    def __getitem__(self, index):\n        path = self.paths[index]\n        img = Image.open(path)\n        return self.transform(img)\n\n\n# stylegan2 classes\n\n\nclass StyleVectorizer(nn.Module):\n    def __init__(self, emb, depth):\n        super().__init__()\n\n        layers = []\n        for i in range(depth):\n            layers.extend([nn.Linear(emb, emb), leaky_relu()])\n\n        self.net = nn.Sequential(*layers)\n\n    def forward(self, x):\n        return self.net(x)\n\n\nclass RGBBlock(nn.Module):\n    def __init__(self, latent_dim, input_channel, upsample, rgba=False):\n        super().__init__()\n        self.input_channel = input_channel\n        self.to_style = nn.Linear(latent_dim, input_channel)\n\n        out_filters = 3 if not rgba else 4\n        self.conv = Conv2DMod(input_channel, out_filters, 1, demod=False)\n\n        self.upsample = nn.Upsample(scale_factor=2, mode=\"bilinear\", align_corners=False) if upsample else None\n\n    def forward(self, x, prev_rgb, istyle):\n        b, c, h, w = x.shape\n        style = self.to_style(istyle)\n        x = self.conv(x, style)\n\n        if prev_rgb is not None:\n            x = x + prev_rgb\n\n        if self.upsample is not None:\n            x = self.upsample(x)\n\n        return x\n\n\nclass Conv2DMod(nn.Module):\n    def __init__(self, in_chan, out_chan, kernel, demod=True, stride=1, dilation=1, **kwargs):\n        super().__init__()\n        self.filters = out_chan\n        self.demod = demod\n        self.kernel = kernel\n        self.stride = stride\n        self.dilation = dilation\n        self.weight = nn.Parameter(torch.randn((out_chan, in_chan, kernel, kernel)))\n        nn.init.kaiming_normal_(self.weight, a=0, mode=\"fan_in\", nonlinearity=\"leaky_relu\")\n\n    def _get_same_padding(self, size, kernel, dilation, stride):\n        return ((size - 1) * (stride - 1) + dilation * (kernel - 1)) // 2\n\n    def forward(self, x, y):\n        b, c, h, w = x.shape\n\n        w1 = y[:, None, :, None, None]\n        w2 = self.weight[None, :, :, :, :]\n        weights = w2 * (w1 + 1)\n\n        if self.demod:\n            d = torch.rsqrt((weights ** 2).sum(dim=(2, 3, 4), keepdim=True) + EPS)\n            weights = weights * d\n\n        x = x.reshape(1, -1, h, w)\n\n        _, _, *ws = weights.shape\n        weights = weights.reshape(b * self.filters, *ws)\n\n        padding = self._get_same_padding(h, self.kernel, self.dilation, self.stride)\n        x = F.conv2d(x, weights, padding=padding, groups=b)\n\n        x = x.reshape(-1, self.filters, h, w)\n        return x\n\n\nclass GeneratorBlock(nn.Module):\n    def __init__(self, latent_dim, input_channels, filters, upsample=True, upsample_rgb=True, rgba=False):\n        super().__init__()\n        self.upsample = nn.Upsample(scale_factor=2, mode=\"bilinear\", align_corners=False) if upsample else None\n\n        self.to_style1 = nn.Linear(latent_dim, input_channels)\n        self.to_noise1 = nn.Linear(1, filters)\n        self.conv1 = Conv2DMod(input_channels, filters, 3)\n\n        self.to_style2 = nn.Linear(latent_dim, filters)\n        self.to_noise2 = nn.Linear(1, filters)\n        self.conv2 = Conv2DMod(filters, filters, 3)\n\n        self.activation = leaky_relu()\n        self.to_rgb = RGBBlock(latent_dim, filters, upsample_rgb, rgba)\n\n    def forward(self, x, prev_rgb, istyle, inoise):\n        if self.upsample is not None:\n            x = self.upsample(x)\n\n        inoise = inoise[:, : x.shape[2], : x.shape[3], :]\n        noise1 = self.to_noise1(inoise).permute((0, 3, 2, 1))\n        noise2 = self.to_noise2(inoise).permute((0, 3, 2, 1))\n\n        style1 = self.to_style1(istyle)\n        x = self.conv1(x, style1)\n        x = self.activation(x + noise1)\n\n        style2 = self.to_style2(istyle)\n        x = self.conv2(x, style2)\n        x = self.activation(x + noise2)\n\n        rgb = self.to_rgb(x, prev_rgb, istyle)\n        return x, rgb\n\n\nclass DiscriminatorBlock(nn.Module):\n    def __init__(self, input_channels, filters, downsample=True):\n        super().__init__()\n        self.conv_res = nn.Conv2d(input_channels, filters, 1)\n\n        self.net = nn.Sequential(\n            nn.Conv2d(input_channels, filters, 3, padding=1),\n            leaky_relu(),\n            nn.Conv2d(filters, filters, 3, padding=1),\n            leaky_relu(),\n        )\n\n        self.downsample = nn.Conv2d(filters, filters, 3, padding=1, stride=2) if downsample else None\n\n    def forward(self, x):\n        res = self.conv_res(x)\n        x = self.net(x)\n        x = x + res\n        if self.downsample is not None:\n            x = self.downsample(x)\n        return x\n\n\nclass Generator(nn.Module):\n    def __init__(self, image_size, latent_dim, network_capacity=16, transparent=False, attn_layers=[]):\n        super().__init__()\n        self.image_size = image_size\n        self.latent_dim = latent_dim\n        self.num_layers = int(log2(image_size) - 1)\n\n        init_channels = 4 * network_capacity\n        self.initial_block = nn.Parameter(torch.randn((init_channels, 4, 4)))\n        filters = [init_channels] + [network_capacity * (2 ** (i + 1)) for i in range(self.num_layers)][::-1]\n        in_out_pairs = zip(filters[0:-1], filters[1:])\n\n        self.blocks = nn.ModuleList([])\n        self.attns = nn.ModuleList([])\n\n        for ind, (in_chan, out_chan) in enumerate(in_out_pairs):\n            not_first = ind != 0\n            not_last = ind != (self.num_layers - 1)\n            num_layer = self.num_layers - ind\n\n            attn_fn = (\n                nn.Sequential(*[Residual(Rezero(ImageLinearAttention(in_chan))) for _ in range(2)])\n                if num_layer in attn_layers\n                else None\n            )\n\n            self.attns.append(attn_fn)\n\n            block = GeneratorBlock(\n                latent_dim, in_chan, out_chan, upsample=not_first, upsample_rgb=not_last, rgba=transparent\n            )\n            self.blocks.append(block)\n\n    def forward(self, styles, input_noise):\n        batch_size = styles.shape[0]\n        image_size = self.image_size\n        x = self.initial_block.expand(batch_size, -1, -1, -1)\n        styles = styles.transpose(0, 1)\n\n        rgb = None\n        for style, block, attn in zip(styles, self.blocks, self.attns):\n            if attn is not None:\n                x = attn(x)\n            x, rgb = block(x, rgb, style, input_noise)\n\n        return rgb\n\n\nclass Discriminator(nn.Module):\n    def __init__(\n        self, image_size, network_capacity=16, fq_layers=[], fq_dict_size=256, attn_layers=[], transparent=False\n    ):\n        super().__init__()\n        num_layers = int(log2(image_size) - 1)\n        num_init_filters = 3 if not transparent else 4\n\n        blocks = []\n        filters = [num_init_filters] + [(network_capacity) * (2 ** i) for i in range(num_layers + 1)]\n        chan_in_out = list(zip(filters[0:-1], filters[1:]))\n\n        blocks = []\n        quantize_blocks = []\n        attn_blocks = []\n\n        for ind, (in_chan, out_chan) in enumerate(chan_in_out):\n            num_layer = ind + 1\n            is_not_last = ind != (len(chan_in_out) - 1)\n\n            block = DiscriminatorBlock(in_chan, out_chan, downsample=is_not_last)\n            blocks.append(block)\n\n            attn_fn = (\n                nn.Sequential(*[Residual(Rezero(ImageLinearAttention(out_chan))) for _ in range(2)])\n                if num_layer in attn_layers\n                else None\n            )\n\n            attn_blocks.append(attn_fn)\n\n            quantize_fn = PermuteToFrom(VectorQuantize(out_chan, fq_dict_size)) if num_layer in fq_layers else None\n            quantize_blocks.append(quantize_fn)\n\n        self.blocks = nn.ModuleList(blocks)\n        self.attn_blocks = nn.ModuleList(attn_blocks)\n        self.quantize_blocks = nn.ModuleList(quantize_blocks)\n\n        latent_dim = 2 * 2 * filters[-1]\n\n        self.flatten = Flatten()\n        self.to_logit = nn.Linear(latent_dim, 1)\n\n    def forward(self, x):\n        b, *_ = x.shape\n\n        quantize_loss = torch.zeros(1).to(x)\n\n        for (block, attn_block, q_block) in zip(self.blocks, self.attn_blocks, self.quantize_blocks):\n            x = block(x)\n\n            if attn_block is not None:\n                x = attn_block(x)\n\n            if q_block is not None:\n                x, loss = q_block(x)\n                quantize_loss += loss\n\n        x = self.flatten(x)\n        x = self.to_logit(x)\n        return x.squeeze(), quantize_loss\n\n\nclass StyleGAN2(nn.Module):\n    def __init__(\n        self,\n        image_size,\n        latent_dim=512,\n        style_depth=8,\n        network_capacity=16,\n        transparent=False,\n        fp16=False,\n        cl_reg=False,\n        augment_fn=None,\n        steps=1,\n        lr=1e-4,\n        fq_layers=[],\n        fq_dict_size=256,\n        attn_layers=[],\n    ):\n        super().__init__()\n        self.lr = lr\n        self.steps = steps\n        self.ema_updater = EMA(0.995)\n\n        self.S = StyleVectorizer(latent_dim, style_depth)\n        self.G = Generator(image_size, latent_dim, network_capacity, transparent=transparent, attn_layers=attn_layers)\n        self.D = Discriminator(\n            image_size,\n            network_capacity,\n            fq_layers=fq_layers,\n            fq_dict_size=fq_dict_size,\n            attn_layers=attn_layers,\n            transparent=transparent,\n        )\n\n        self.SE = StyleVectorizer(latent_dim, style_depth)\n        self.GE = Generator(image_size, latent_dim, network_capacity, transparent=transparent, attn_layers=attn_layers)\n\n        set_requires_grad(self.SE, False)\n        set_requires_grad(self.GE, False)\n\n        generator_params = list(self.G.parameters()) + list(self.S.parameters())\n        self.G_opt = DiffGrad(generator_params, lr=self.lr, betas=(0.5, 0.9))\n        self.D_opt = DiffGrad(self.D.parameters(), lr=self.lr, betas=(0.5, 0.9))\n\n        self._init_weights()\n        self.reset_parameter_averaging()\n\n        self.cuda()\n\n        if fp16:\n            (self.S, self.G, self.D, self.SE, self.GE), (self.G_opt, self.D_opt) = amp.initialize(\n                [self.S, self.G, self.D, self.SE, self.GE], [self.G_opt, self.D_opt], opt_level=\"O2\"\n            )\n\n        # experimental contrastive loss discriminator regularization\n        if augment_fn is not None:\n            self.augment_fn = augment_fn\n        else:\n            self.augment_fn = nn.Sequential(\n                nn.ReflectionPad2d(int((sqrt(2) - 1) * image_size / 4)),\n                RandomApply(augs.ColorJitter(0.8, 0.8, 0.8, 0.2), p=0.7),\n                augs.RandomGrayscale(p=0.2),\n                augs.RandomHorizontalFlip(),\n                RandomApply(augs.RandomAffine(degrees=0, translate=(0.25, 0.25), shear=(15, 15)), p=0.3),\n                RandomApply(\n                    nn.Sequential(augs.RandomRotation(180), augs.CenterCrop(size=(image_size, image_size))), p=0.2\n                ),\n                augs.RandomResizedCrop(size=(image_size, image_size)),\n                RandomApply(filters.GaussianBlur2d((3, 3), (1.5, 1.5)), p=0.1),\n                RandomApply(augs.RandomErasing(), p=0.1),\n            )\n\n        self.D_cl = (\n            ContrastiveLearner(self.D, image_size, augment_fn=self.augment_fn, fp16=fp16, hidden_layer=\"flatten\")\n            if cl_reg\n            else None\n        )\n\n        # self.S, self.G, self.D, self.SE, self.GE = (\n        #     nn.DataParallel(self.S),\n        #     nn.DataParallel(self.G),\n        #     nn.DataParallel(self.D),\n        #     nn.DataParallel(self.SE),\n        #     nn.DataParallel(self.GE),\n        # )\n\n    def _init_weights(self):\n        for m in self.modules():\n            if type(m) in {nn.Conv2d, nn.Linear}:\n                nn.init.kaiming_normal_(m.weight, a=0, mode=\"fan_in\", nonlinearity=\"leaky_relu\")\n\n        for block in self.G.blocks:\n            nn.init.zeros_(block.to_noise1.weight)\n            nn.init.zeros_(block.to_noise2.weight)\n            nn.init.zeros_(block.to_noise1.bias)\n            nn.init.zeros_(block.to_noise2.bias)\n\n    def EMA(self):\n        def update_moving_average(ma_model, current_model):\n            for current_params, ma_params in zip(current_model.parameters(), ma_model.parameters()):\n                old_weight, up_weight = ma_params.data, current_params.data\n                ma_params.data = self.ema_updater.update_average(old_weight, up_weight)\n\n        update_moving_average(self.SE, self.S)\n        update_moving_average(self.GE, self.G)\n\n    def reset_parameter_averaging(self):\n        self.SE.load_state_dict(self.S.state_dict())\n        self.GE.load_state_dict(self.G.state_dict())\n\n    def forward(self, x):\n        return x\n\n\nclass Trainer:\n    def __init__(\n        self,\n        name,\n        results_dir,\n        models_dir,\n        image_size,\n        network_capacity,\n        transparent=False,\n        batch_size=4,\n        mixed_prob=0.9,\n        gradient_accumulate_every=1,\n        lr=2e-4,\n        num_workers=None,\n        save_every=1000,\n        trunc_psi=0.6,\n        fp16=False,\n        cl_reg=False,\n        fq_layers=[],\n        fq_dict_size=256,\n        attn_layers=[],\n        fid_n_sample=5000,\n        ppl_n_sample=2500,\n        *args,\n        **kwargs,\n    ):\n        self.GAN_params = [args, kwargs]\n        self.GAN = None\n\n        self.name = name\n        self.results_dir = Path(results_dir)\n        self.models_dir = Path(models_dir)\n        self.config_path = self.models_dir / name / \".config.json\"\n\n        assert log2(image_size).is_integer(), \"image size must be a power of 2 (64, 128, 256, 512, 1024)\"\n        self.image_size = image_size\n        self.network_capacity = network_capacity\n        self.transparent = transparent\n        self.fq_layers = cast_list(fq_layers)\n        self.fq_dict_size = fq_dict_size\n\n        self.attn_layers = cast_list(attn_layers)\n\n        self.lr = lr\n        self.batch_size = batch_size\n        self.num_workers = num_workers\n        self.mixed_prob = mixed_prob\n\n        self.save_every = save_every\n        self.steps = 0\n\n        self.av = None\n        self.trunc_psi = trunc_psi\n\n        self.pl_mean = 0\n\n        self.fid_n_sample = fid_n_sample\n        self.ppl_n_sample = ppl_n_sample\n\n        self.gradient_accumulate_every = gradient_accumulate_every\n\n        assert not fp16 or fp16 and APEX_AVAILABLE, \"Apex is not available for you to use mixed precision training\"\n        self.fp16 = fp16\n\n        self.cl_reg = cl_reg\n\n        self.d_loss = 0\n        self.g_loss = 0\n        self.last_gp_loss = 0\n        self.last_cr_loss = 0\n        self.q_loss = 0\n\n        self.pl_length_ma = EMA(0.99)\n        self.init_folders()\n\n        self.loader = None\n\n    def init_GAN(self):\n        args, kwargs = self.GAN_params\n        self.GAN = StyleGAN2(\n            lr=self.lr,\n            image_size=self.image_size,\n            network_capacity=self.network_capacity,\n            transparent=self.transparent,\n            fq_layers=self.fq_layers,\n            fq_dict_size=self.fq_dict_size,\n            attn_layers=self.attn_layers,\n            fp16=self.fp16,\n            cl_reg=self.cl_reg,\n            *args,\n            **kwargs,\n        )\n\n    def write_config(self):\n        self.config_path.write_text(json.dumps(self.config()))\n\n    def load_config(self):\n        config = self.config() if not self.config_path.exists() else json.loads(self.config_path.read_text())\n        self.image_size = config[\"image_size\"]\n        self.network_capacity = config[\"network_capacity\"]\n        self.transparent = config[\"transparent\"]\n        self.fq_layers = config[\"fq_layers\"]\n        self.fq_dict_size = config[\"fq_dict_size\"]\n        self.attn_layers = config.pop(\"attn_layers\", [])\n        del self.GAN\n        self.init_GAN()\n\n    def config(self):\n        return {\n            \"image_size\": self.image_size,\n            \"network_capacity\": self.network_capacity,\n            \"transparent\": self.transparent,\n            \"fq_layers\": self.fq_layers,\n            \"fq_dict_size\": self.fq_dict_size,\n            \"attn_layers\": self.attn_layers,\n        }\n\n    def set_data_src(self, folder):\n        self.dataset = Dataset(folder, self.image_size, transparent=self.transparent)\n        self.loader = cycle(\n            data.DataLoader(\n                self.dataset,\n                num_workers=default(self.num_workers, num_cores),\n                batch_size=self.batch_size,\n                drop_last=True,\n                shuffle=True,\n                pin_memory=True,\n            )\n        )\n        validation.get_dataset_inception_features(self.loader, self.name, self.image_size)\n\n    def train(self):\n        assert (\n            self.loader is not None\n        ), \"You must first initialize the data source with `.set_data_src(<folder of images>)`\"\n\n        if self.GAN is None:\n            self.init_GAN()\n\n        self.GAN.train()\n        total_disc_loss = torch.tensor(0.0).cuda()\n        total_gen_loss = torch.tensor(0.0).cuda()\n\n        batch_size = self.batch_size\n\n        image_size = self.GAN.G.image_size\n        latent_dim = self.GAN.G.latent_dim\n        num_layers = self.GAN.G.num_layers\n\n        apply_gradient_penalty = self.steps % 4 == 0\n        apply_path_penalty = self.steps % 32 == 0\n        apply_cl_reg_to_generated = self.steps > 20000\n\n        log_dict = {\"Divergence\": 0, \"Quantize\": 0, \"Generator\": 0}\n        if apply_gradient_penalty:\n            log_dict[\"R1\"] = 0\n        if apply_path_penalty:\n            log_dict[\"Path Length\"] = 0\n\n        backwards = partial(loss_backwards, self.fp16)\n\n        if self.GAN.D_cl is not None:\n            self.GAN.D_opt.zero_grad()\n\n            if apply_cl_reg_to_generated:\n                for i in range(self.gradient_accumulate_every):\n                    get_latents_fn = mixed_list if random() < self.mixed_prob else noise_list\n                    style = get_latents_fn(batch_size, num_layers, latent_dim)\n                    noise = image_noise(batch_size, image_size)\n\n                    w_space = latent_to_w(self.GAN.S, style)\n                    w_styles = styles_def_to_tensor(w_space)\n\n                    generated_images = self.GAN.G(w_styles, noise)\n                    self.GAN.D_cl(generated_images.clone().detach(), accumulate=True)\n\n            for i in range(self.gradient_accumulate_every):\n                image_batch = next(self.loader).cuda()\n                self.GAN.D_cl(image_batch, accumulate=True)\n\n            loss = self.GAN.D_cl.calculate_loss()\n            self.last_cr_loss = loss.clone().detach().item()\n            log_dict[\"Consistency\"] = self.last_cr_loss\n            backwards(loss, self.GAN.D_opt)\n\n            self.GAN.D_opt.step()\n\n        # train discriminator\n\n        avg_pl_length = self.pl_mean\n        self.GAN.D_opt.zero_grad()\n\n        for i in range(self.gradient_accumulate_every):\n            get_latents_fn = mixed_list if random() < self.mixed_prob else noise_list\n            style = get_latents_fn(batch_size, num_layers, latent_dim)\n            noise = image_noise(batch_size, image_size)\n\n            w_space = latent_to_w(self.GAN.S, style)\n            w_styles = styles_def_to_tensor(w_space)\n\n            generated_images = self.GAN.G(w_styles, noise)\n            fake_output, fake_q_loss = self.GAN.D(generated_images.clone().detach())\n\n            image_batch = next(self.loader).cuda()\n            image_batch.requires_grad_()\n            real_output, real_q_loss = self.GAN.D(image_batch)\n\n            divergence = (F.relu(1 + real_output) + F.relu(1 - fake_output)).mean()\n            disc_loss = divergence\n            log_dict[\"Divergence\"] += divergence / self.gradient_accumulate_every\n\n            quantize_loss = (fake_q_loss + real_q_loss).mean()\n            self.q_loss = float(quantize_loss.detach().item())\n            log_dict[\"Quantize\"] += self.q_loss / self.gradient_accumulate_every\n\n            disc_loss = disc_loss + quantize_loss\n\n            if apply_gradient_penalty:\n                gp = gradient_penalty(image_batch, real_output)\n                self.last_gp_loss = gp.clone().detach().item()\n                disc_loss = disc_loss + gp\n                log_dict[\"R1\"] += gp / self.gradient_accumulate_every\n\n            disc_loss = disc_loss / self.gradient_accumulate_every\n            disc_loss.register_hook(raise_if_nan)\n            backwards(disc_loss, self.GAN.D_opt)\n\n            total_disc_loss += divergence.detach().item() / self.gradient_accumulate_every\n\n        self.d_loss = float(total_disc_loss)\n        log_dict[\"Discriminator\"] = self.d_loss\n        self.GAN.D_opt.step()\n\n        # train generator\n\n        self.GAN.G_opt.zero_grad()\n        for i in range(self.gradient_accumulate_every):\n            style = get_latents_fn(batch_size, num_layers, latent_dim)\n            noise = image_noise(batch_size, image_size)\n\n            w_space = latent_to_w(self.GAN.S, style)\n            w_styles = styles_def_to_tensor(w_space)\n\n            generated_images = self.GAN.G(w_styles, noise)\n            fake_output, _ = self.GAN.D(generated_images)\n            loss = fake_output.mean()\n            gen_loss = loss\n            log_dict[\"Generator\"] += gen_loss / self.gradient_accumulate_every\n\n            if apply_path_penalty:\n                std = 0.1 / (w_styles.std(dim=0, keepdim=True) + EPS)\n                w_styles_2 = w_styles + torch.randn(w_styles.shape).cuda() / (std + EPS)\n                pl_images = self.GAN.G(w_styles_2, noise)\n                pl_lengths = ((pl_images - generated_images) ** 2).mean(dim=(1, 2, 3))\n                avg_pl_length = np.mean(pl_lengths.detach().cpu().numpy())\n\n                if not is_empty(self.pl_mean):\n                    pl_loss = ((pl_lengths - self.pl_mean) ** 2).mean()\n                    log_dict[\"Path Length\"] += pl_loss / self.gradient_accumulate_every\n                    if not torch.isnan(pl_loss):\n                        gen_loss = gen_loss + pl_loss\n\n            gen_loss = gen_loss / self.gradient_accumulate_every\n            gen_loss.register_hook(raise_if_nan)\n            backwards(gen_loss, self.GAN.G_opt)\n\n            total_gen_loss += loss.detach().item() / self.gradient_accumulate_every\n\n        self.g_loss = float(total_gen_loss)\n        self.GAN.G_opt.step()\n\n        # calculate moving averages\n        if apply_path_penalty and not np.isnan(avg_pl_length):\n            self.pl_mean = self.pl_length_ma.update_average(self.pl_mean, avg_pl_length)\n            log_dict[\"Mean Path Length\"] = self.pl_mean\n\n        if self.steps % 10 == 0 and self.steps > 20000:\n            self.GAN.EMA()\n\n        if self.steps <= 25000 and self.steps % 1000 == 2:\n            self.GAN.reset_parameter_averaging()\n\n        # save from NaN errors\n        checkpoint_num = floor(self.steps / self.save_every)\n\n        if any(torch.isnan(l) for l in (total_gen_loss, total_disc_loss)):\n            print(f\"NaN detected for generator or discriminator. Loading from checkpoint #{checkpoint_num}\")\n            self.load(checkpoint_num)\n            raise NanException\n\n        # periodically save results\n        if self.steps % self.save_every == 0:\n            self.save(checkpoint_num)\n\n        if self.steps % 1000 == 0 or (self.steps % 100 == 0 and self.steps < 2500):\n            self.evaluate(floor(self.steps / 1000))\n\n        if self.steps % 1000 == 0:\n            start_time = time.time()\n\n            PBAR.set_description((f\"Calculating FID...\"))\n            fid, density, coverage = self.calculate_fid()\n            log_dict[\"Evaluation/FID\"] = fid\n            log_dict[\"Evaluation/Density\"] = density\n            log_dict[\"Evaluation/Coverage\"] = coverage\n\n            PBAR.set_description((f\"Calculating PPL...\"))\n            ppl = self.calculate_ppl()\n\n            PBAR.set_description(\n                (\n                    f\"FID: {fid:.4f}; Density: {density:.4f}; Coverage: {coverage:.4f}; PPL: {ppl:.4f} in {time.time() - start_time:.1f}s\"\n                )\n            )\n            log_dict[\"Evaluation/PPL\"] = ppl\n\n        wandb.log(log_dict)\n\n        self.steps += 1\n        self.av = None\n\n    @torch.no_grad()\n    def evaluate(self, num=0, num_image_tiles=8):\n        self.GAN.eval()\n        ext = \"jpg\" if not self.transparent else \"png\"\n        num_rows = num_image_tiles\n\n        latent_dim = self.GAN.G.latent_dim\n        image_size = self.GAN.G.image_size\n        num_layers = self.GAN.G.num_layers\n\n        # latents and noise\n        latents = noise_list(num_rows ** 2, num_layers, latent_dim)\n        n = image_noise(num_rows ** 2, image_size)\n\n        # regular\n        generated_images = self.generate_truncated(self.GAN.S, self.GAN.G, latents, n, trunc_psi=self.trunc_psi)\n        grid = torchvision.utils.make_grid(generated_images, nrow=num_rows)\n        wandb.log({\"Generated Images\": [wandb.Image(grid, caption=f\"Step {num}\")]})\n\n        # moving averages\n        generated_images = self.generate_truncated(self.GAN.SE, self.GAN.GE, latents, n, trunc_psi=self.trunc_psi)\n        grid = torchvision.utils.make_grid(generated_images, nrow=num_rows)\n        wandb.log({\"Generated Images EMA\": [wandb.Image(grid, caption=f\"Step {num}\")]})\n\n        # mixing regularities\n        def tile(a, dim, n_tile):\n            init_dim = a.size(dim)\n            repeat_idx = [1] * a.dim()\n            repeat_idx[dim] = n_tile\n            a = a.repeat(*(repeat_idx))\n            order_index = torch.LongTensor(\n                np.concatenate([init_dim * np.arange(n_tile) + i for i in range(init_dim)])\n            ).cuda()\n            return torch.index_select(a, dim, order_index)\n\n        nn = noise(num_rows, latent_dim)\n        tmp1 = tile(nn, 0, num_rows)\n        tmp2 = nn.repeat(num_rows, 1)\n\n        tt = int(num_layers / 2)\n        mixed_latents = [(tmp1, tt), (tmp2, num_layers - tt)]\n\n        generated_images = self.generate_truncated(self.GAN.SE, self.GAN.GE, mixed_latents, n, trunc_psi=self.trunc_psi)\n        grid = torchvision.utils.make_grid(generated_images, nrow=num_rows)\n        wandb.log({\"Style Mixing\": [wandb.Image(grid, caption=f\"Step {num}\")]})\n\n    @torch.no_grad()\n    def calculate_fid(self):\n        self.GAN.eval()\n\n        inception = InceptionV3([3], normalize_input=False, init_weights=False)\n        inception = inception.eval().to(next(self.GAN.parameters()).device)\n\n        latent_dim = self.GAN.G.latent_dim\n        image_size = self.GAN.G.image_size\n        num_layers = self.GAN.G.num_layers\n\n        features = []\n        for _ in range(floor(self.fid_n_sample / self.batch_size) + 1):\n            latents = noise_list(self.batch_size, num_layers, latent_dim)\n            n = image_noise(self.batch_size, image_size)\n            imgs = self.generate_truncated(self.GAN.SE, self.GAN.GE, latents, n, trunc_psi=self.trunc_psi)\n            feat = inception(imgs)[0].view(imgs.shape[0], -1)\n            features.append(feat.to(\"cpu\"))\n        features = torch.cat(features, 0).numpy()\n\n        del inception\n\n        sample_mean = np.mean(features, 0)\n        sample_cov = np.cov(features, rowvar=False)\n\n        with open(f\"inception_{self.name}_stats.pkl\", \"rb\") as f:\n            embeds = pickle.load(f)\n            real_mean = embeds[\"mean\"]\n            real_cov = embeds[\"cov\"]\n\n        cov_sqrt, _ = linalg.sqrtm(sample_cov @ real_cov, disp=False)\n        if not np.isfinite(cov_sqrt).all():\n            print(\"product of cov matrices is singular\")\n            offset = np.eye(sample_cov.shape[0]) * 1e-6\n            cov_sqrt = linalg.sqrtm((sample_cov + offset) @ (real_cov + offset))\n        if np.iscomplexobj(cov_sqrt):\n            if not np.allclose(np.diagonal(cov_sqrt).imag, 0, atol=1e-3):\n                raise ValueError(f\"Imaginary component {np.max(np.abs(cov_sqrt.imag))}\")\n            cov_sqrt = cov_sqrt.real\n\n        mean_diff = sample_mean - real_mean\n        mean_norm = mean_diff @ mean_diff\n        trace = np.trace(sample_cov) + np.trace(real_cov) - 2 * np.trace(cov_sqrt)\n        inception_distance = mean_norm + trace\n\n        with open(f\"inception_{self.name}_features.pkl\", \"rb\") as f:\n            embeds = pickle.load(f)\n            real_feats = embeds[\"features\"]\n        _, _, density, coverage = validation.prdc(real_feats[:80000], features[:80000])\n\n        return inception_distance, density, coverage\n\n    @torch.no_grad()\n    def calculate_ppl(self):\n        self.GAN.eval()\n        latent_dim = self.GAN.G.latent_dim\n        image_size = self.GAN.G.image_size\n        num_layers = self.GAN.G.num_layers\n        eps = 1e-4\n\n        def lerp(a, b, t):\n            return a + (b - a) * t\n\n        percept = lpips.PerceptualLoss(\n            model=\"net-lin\", net=\"vgg\", use_gpu=True, gpu_ids=[next(self.GAN.parameters()).device.index]\n        )\n\n        distances = []\n        for _ in range(floor(self.fid_n_sample / self.batch_size) + 1):\n            noise = image_noise(self.batch_size * 2, image_size)\n            inputs = noise_list(self.batch_size * 2, num_layers, latent_dim)\n            lerp_t = torch.rand(self.batch_size).cuda()\n\n            # print(lerp_t.shape)\n\n            w_space = []\n            for tensor, num_layers in inputs:\n                av = torch.from_numpy(self.av).cuda()\n                tmp = self.trunc_psi * (self.GAN.SE(tensor) - av) + av\n                w_space.append((tmp, num_layers))\n            latent = styles_def_to_tensor(w_space)\n\n            latent_t0, latent_t1 = latent[::2], latent[1::2]\n            latent_e0 = lerp(latent_t0, latent_t1, lerp_t[:, None, None])\n            latent_e1 = lerp(latent_t0, latent_t1, lerp_t[:, None, None] + eps)\n            latent_e = torch.stack([latent_e0, latent_e1], 1).view(*latent.shape)\n\n            image = evaluate_in_chunks(self.batch_size, self.GAN.GE, latent_e, noise)\n\n            factor = image.shape[2] // 256\n            if factor > 1:\n                image = F.interpolate(image, size=(256, 256), mode=\"bilinear\", align_corners=False)\n\n            dist = percept(image[::2], image[1::2]).view(image.shape[0] // 2) / (eps ** 2)\n            distances.append(dist.to(\"cpu\").numpy())\n        distances = np.concatenate(distances, 0)\n\n        lo = np.percentile(distances, 1, interpolation=\"lower\")\n        hi = np.percentile(distances, 99, interpolation=\"higher\")\n        filtered_dist = np.extract(np.logical_and(lo <= distances, distances <= hi), distances)\n        path_length = filtered_dist.mean()\n\n        del percept, inputs, lerp_t, image, dist\n\n        return path_length\n\n    @torch.no_grad()\n    def generate_truncated(self, S, G, style, noi, trunc_psi=0.75, num_image_tiles=8):\n        latent_dim = G.latent_dim\n\n        if self.av is None:\n            z = noise(2 ** 14, latent_dim)\n            samples = evaluate_in_chunks(self.batch_size, S, z).cpu().numpy()\n            self.av = np.mean(samples, axis=0)\n            self.av = np.expand_dims(self.av, axis=0)\n\n        w_space = []\n        for tensor, num_layers in style:\n            tmp = S(tensor)\n            av_torch = torch.from_numpy(self.av).cuda()\n            tmp = trunc_psi * (tmp - av_torch) + av_torch\n            w_space.append((tmp, num_layers))\n\n        w_styles = styles_def_to_tensor(w_space)\n        generated_images = evaluate_in_chunks(self.batch_size, G, w_styles, noi)\n        return generated_images.clamp_(0.0, 1.0)\n\n    def print_log(self):\n        print(\n            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}\"\n        )\n\n    def model_name(self, num):\n        return str(self.models_dir / self.name / f\"model_{wandb.run.dir.split('/')[-1].split('-')[-1]}_{num}.pt\")\n\n    def init_folders(self):\n        (self.results_dir / self.name).mkdir(parents=True, exist_ok=True)\n        (self.models_dir / self.name).mkdir(parents=True, exist_ok=True)\n\n    def clear(self):\n        rmtree(f\"./models/{self.name}\", True)\n        rmtree(f\"./results/{self.name}\", True)\n        rmtree(str(self.config_path), True)\n        self.init_folders()\n\n    def save(self, num):\n        torch.save(self.GAN.state_dict(), self.model_name(num))\n        self.write_config()\n\n    def load(self, num=-1):\n        self.load_config()\n\n        name = num\n        if num == -1:\n            file_paths = [p for p in Path(self.models_dir / self.name).glob(\"model_*.pt\")]\n            saved_nums = sorted(map(lambda x: int(x.stem.split(\"_\")[-1]), file_paths))\n            if len(saved_nums) == 0:\n                return\n            name = saved_nums[-1]\n            print(f\"continuing from previous epoch - {name}\")\n        self.steps = name * self.save_every\n        self.GAN.load_state_dict(torch.load(self.model_name(name)))\n\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser()\n    parser.add_argument(\"data\", type=str)\n    parser.add_argument(\"name\", type=str)\n    parser.add_argument(\"--results_dir\", type=str, default=\"/home/hans/neurout/\")\n    parser.add_argument(\"--models_dir\", type=str, default=\"/home/hans/modelzoo/maua-sg2/\")\n    parser.add_argument(\"--new\", type=bool, default=False)\n    parser.add_argument(\"--load_from\", type=str, default=-1)\n    parser.add_argument(\"--image_size\", type=int, default=256)\n    parser.add_argument(\"--network_capacity\", type=int, default=16)\n    parser.add_argument(\"--transparent\", type=bool, default=False)\n    parser.add_argument(\"--batch_size\", type=int, default=2)\n    parser.add_argument(\"--gradient_accumulate_every\", type=int, default=12)\n    parser.add_argument(\"--num_train_steps\", type=int, default=150000)\n    parser.add_argument(\"--learning_rate\", type=float, default=2e-4)\n    parser.add_argument(\"--num_workers\", type=int, default=None)\n    parser.add_argument(\"--save_every\", type=int, default=1000)\n    parser.add_argument(\"--generate\", type=bool, default=False)\n    parser.add_argument(\"--num_image_tiles\", type=int, default=8)\n    parser.add_argument(\"--trunc_psi\", type=float, default=1)\n    parser.add_argument(\"--fp16\", type=bool, default=False)\n    parser.add_argument(\"--cl_reg\", type=bool, default=True)\n    parser.add_argument(\"--fq_layers\", default=[])\n    parser.add_argument(\"--fq_dict_size\", default=256)\n    parser.add_argument(\"--attn_layers\", default=[])\n    args = parser.parse_args()\n\n    wandb.init(project=f\"maua-stylegan\", name=\"lucidrains-\" + args.name)\n\n    model = Trainer(\n        args.name,\n        args.results_dir,\n        args.models_dir,\n        batch_size=args.batch_size,\n        gradient_accumulate_every=args.gradient_accumulate_every,\n        image_size=args.image_size,\n        network_capacity=args.network_capacity,\n        transparent=args.transparent,\n        lr=args.learning_rate,\n        num_workers=args.num_workers,\n        save_every=args.save_every,\n        trunc_psi=args.trunc_psi,\n        fp16=args.fp16,\n        cl_reg=args.cl_reg,\n        fq_layers=args.fq_layers,\n        fq_dict_size=args.fq_dict_size,\n        attn_layers=args.attn_layers,\n    )\n\n    if not args.new:\n        model.load(args.load_from)\n    else:\n        model.clear()\n\n    if args.generate:\n        now = datetime.now()\n        timestamp = now.strftime(\"%m-%d-%Y_%H-%M-%S\")\n        samples_name = f\"generated-{timestamp}\"\n        model.evaluate(samples_name, args.num_image_tiles)\n        print(f\"sample images generated at {args.results_dir}/{args.name}/{args.samples_name}\")\n        exit()\n\n    model.set_data_src(args.data)\n\n    PBAR = tqdm(range(args.num_train_steps - model.steps), mininterval=10.0, desc=f\"{args.name}<{args.data}>\")\n    for _ in PBAR:\n        retry_call(model.train, tries=3, exceptions=NanException)\n        if _ % 50 == 0:\n            model.print_log()\n"
  },
  {
    "path": "models/autoencoder.py",
    "content": "import os\nimport sys\nfrom copy import copy\n\nimport torch as th\nimport torch.nn.functional as F\n\nsys.path.append(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))\n\nfrom op import FusedLeakyReLU\n\n\ndef info(x):\n    print(x.shape, x.min(), x.mean(), x.max())\n\n\nclass PrintShape(th.nn.Module):\n    def __init__(self):\n        super(PrintShape, self).__init__()\n\n    def forward(self, x):\n        print(x.shape)\n        return x\n\n\nclass Flatten(th.nn.Module):\n    def forward(self, x):\n        return x.view(x.size(0), -1)\n\n\nclass UnFlatten(th.nn.Module):\n    def __init__(self, channels, size):\n        super(UnFlatten, self).__init__()\n        self.channels = channels\n        self.size = size\n\n    def forward(self, x):\n        return x.view(x.size(0), self.channels, self.size, self.size)\n\n\nclass LogCoshVAE(th.nn.Module):\n    \"\"\"\n    Adapted from https://github.com/AntixK/PyTorch-VAE\n    See LICENSE_AUTOENCODER\n    \"\"\"\n\n    def __init__(self, in_channels, latent_dim, hidden_dims=None, alpha=10.0, beta=1.0, kld_weight=1):\n        super(LogCoshVAE, self).__init__()\n\n        my_hidden_dims = copy(hidden_dims)\n\n        self.latent_dim = latent_dim\n        self.alpha = alpha\n        self.beta = beta\n        self.kld_weight = kld_weight\n\n        modules = []\n        if my_hidden_dims is None:\n            my_hidden_dims = [32, 64, 128, 256, 512]\n\n        # Build Encoder\n        for h_dim in my_hidden_dims:\n            modules.append(\n                th.nn.Sequential(\n                    th.nn.Conv2d(in_channels, out_channels=h_dim, kernel_size=3, stride=2, padding=1),\n                    th.nn.BatchNorm2d(h_dim),\n                    FusedLeakyReLU(h_dim),\n                )\n            )\n            in_channels = h_dim\n\n        self.encoder = th.nn.Sequential(*modules)\n        self.fc_mu = th.nn.Linear(my_hidden_dims[-1] * 4, latent_dim)\n        self.fc_var = th.nn.Linear(my_hidden_dims[-1] * 4, latent_dim)\n\n        # Build Decoder\n        modules = []\n\n        self.decoder_input = th.nn.Linear(latent_dim, my_hidden_dims[-1] * 4)\n\n        my_hidden_dims.reverse()\n\n        for i in range(len(my_hidden_dims) - 1):\n            modules.append(\n                th.nn.Sequential(\n                    th.nn.Upsample(scale_factor=2, mode=\"bilinear\", align_corners=False),\n                    th.nn.Conv2d(my_hidden_dims[i], my_hidden_dims[i + 1], kernel_size=3, padding=1),\n                    th.nn.BatchNorm2d(my_hidden_dims[i + 1]),\n                    FusedLeakyReLU(my_hidden_dims[i + 1]),\n                )\n            )\n\n        self.decoder = th.nn.Sequential(*modules)\n\n        self.final_layer = th.nn.Sequential(\n            th.nn.Upsample(scale_factor=2, mode=\"bilinear\", align_corners=False),\n            th.nn.Conv2d(my_hidden_dims[-1], my_hidden_dims[-1], kernel_size=3, padding=1),\n            th.nn.BatchNorm2d(my_hidden_dims[-1]),\n            FusedLeakyReLU(my_hidden_dims[-1]),\n            th.nn.Conv2d(my_hidden_dims[-1], out_channels=3, kernel_size=3, padding=1),\n            th.nn.Tanh(),\n        )\n\n    def encode(self, input):\n        result = self.encoder(input)\n        result = th.flatten(result, start_dim=1)\n\n        mu = self.fc_mu(result)\n        log_var = self.fc_var(result)\n\n        return mu, log_var\n\n    def decode(self, z):\n        result = self.decoder_input(z)\n        result = result.view(-1, self.latent_dim, 2, 2)\n        result = self.decoder(result)\n        result = self.final_layer(result)\n        return result\n\n    def reparameterize(self, mu, logvar):\n        std = th.exp(0.5 * logvar)\n        eps = th.randn_like(std)\n        return eps * std + mu\n\n    def forward(self, input):\n        mu, log_var = self.encode(input)\n        z = self.reparameterize(mu, log_var)\n        return self.decode(z), mu, log_var\n\n    def loss(self, real, fake, mu, log_var):\n        t = fake - real\n\n        recons_loss = self.alpha * t + th.log(1.0 + th.exp(-2 * self.alpha * t)) - th.log(2.0 * th.ones((1)))\n        recons_loss = (1.0 / self.alpha) * recons_loss.mean()\n\n        kld_loss = th.mean(-0.5 * th.sum(1 + log_var - mu ** 2 - log_var.exp(), dim=1), dim=0)\n\n        loss = recons_loss + self.beta * self.kld_weight * kld_loss\n\n        return {\"Total\": loss, \"Reconstruction\": recons_loss, \"Kullback Leibler Divergence\": -kld_loss}\n\n\nclass conv2DBatchNormRelu(th.nn.Module):\n    def __init__(\n        self, in_channels, n_filters, k_size, stride, padding, bias=True, dilation=1, with_bn=True,\n    ):\n        super(conv2DBatchNormRelu, self).__init__()\n\n        conv_mod = th.nn.Conv2d(\n            int(in_channels),\n            int(n_filters),\n            kernel_size=k_size,\n            padding=padding,\n            stride=stride,\n            bias=bias,\n            dilation=dilation,\n        )\n\n        if with_bn:\n            self.cbr_unit = th.nn.Sequential(\n                conv_mod, th.nn.BatchNorm2d(int(n_filters)), FusedLeakyReLU(int(n_filters))\n            )\n        else:\n            self.cbr_unit = th.nn.Sequential(conv_mod, FusedLeakyReLU(int(n_filters)))\n\n    def forward(self, inputs):\n        outputs = self.cbr_unit(inputs)\n        return outputs\n\n\nclass segnetDown2(th.nn.Module):\n    def __init__(self, in_size, out_size):\n        super(segnetDown2, self).__init__()\n        self.conv1 = conv2DBatchNormRelu(in_size, out_size, 3, 1, 1)\n        self.conv2 = conv2DBatchNormRelu(out_size, out_size, 3, 1, 1)\n        self.maxpool_with_argmax = th.nn.MaxPool2d(2, 2, return_indices=True)\n\n    def forward(self, inputs):\n        outputs = self.conv1(inputs)\n        outputs = self.conv2(outputs)\n        unpooled_shape = outputs.size()\n        outputs, indices = self.maxpool_with_argmax(outputs)\n        return outputs, indices, unpooled_shape\n\n\nclass segnetDown3(th.nn.Module):\n    def __init__(self, in_size, out_size):\n        super(segnetDown3, self).__init__()\n        self.conv1 = conv2DBatchNormRelu(in_size, out_size, 3, 1, 1)\n        self.conv2 = conv2DBatchNormRelu(out_size, out_size, 3, 1, 1)\n        self.conv3 = conv2DBatchNormRelu(out_size, out_size, 3, 1, 1)\n        self.maxpool_with_argmax = th.nn.MaxPool2d(2, 2, return_indices=True)\n\n    def forward(self, inputs):\n        outputs = self.conv1(inputs)\n        outputs = self.conv2(outputs)\n        outputs = self.conv3(outputs)\n        unpooled_shape = outputs.size()\n        outputs, indices = self.maxpool_with_argmax(outputs)\n        return outputs, indices, unpooled_shape\n\n\nclass segnetUp2(th.nn.Module):\n    def __init__(self, in_size, out_size):\n        super(segnetUp2, self).__init__()\n        self.unpool = th.nn.MaxUnpool2d(2, 2)\n        self.conv1 = conv2DBatchNormRelu(in_size, in_size, 3, 1, 1)\n        self.conv2 = conv2DBatchNormRelu(in_size, out_size, 3, 1, 1)\n\n    def forward(self, inputs, indices, output_shape):\n        outputs = self.unpool(input=inputs, indices=indices, output_size=output_shape)\n        outputs = self.conv1(outputs)\n        outputs = self.conv2(outputs)\n        return outputs\n\n\nclass segnetUp3(th.nn.Module):\n    def __init__(self, in_size, out_size):\n        super(segnetUp3, self).__init__()\n        self.unpool = th.nn.MaxUnpool2d(2, 2)\n        self.conv1 = conv2DBatchNormRelu(in_size, in_size, 3, 1, 1)\n        self.conv2 = conv2DBatchNormRelu(in_size, in_size, 3, 1, 1)\n        self.conv3 = conv2DBatchNormRelu(in_size, out_size, 3, 1, 1)\n\n    def forward(self, inputs, indices, output_shape):\n        outputs = self.unpool(input=inputs, indices=indices, output_size=output_shape)\n        outputs = self.conv1(outputs)\n        outputs = self.conv2(outputs)\n        outputs = self.conv3(outputs)\n        return outputs\n\n\nclass SegNet(th.nn.Module):\n    \"\"\"\n    Adapted from https://github.com/foamliu/Autoencoder\n    See LICENSE_AUTOENCODER\n    \"\"\"\n\n    def __init__(self, in_channels=3):\n        super(SegNet, self).__init__()\n\n        self.down1 = segnetDown2(in_channels, 64)\n        self.down2 = segnetDown2(64, 128)\n        self.down3 = segnetDown3(128, 256)\n        self.down4 = segnetDown3(256, 512)\n        self.down5 = segnetDown3(512, 512)\n\n        self.up5 = segnetUp3(512, 512)\n        self.up4 = segnetUp3(512, 256)\n        self.up3 = segnetUp3(256, 128)\n        self.up2 = segnetUp2(128, 64)\n        self.up1 = segnetUp2(64, in_channels)\n\n    def random_indices(self, shape):\n        batch, channel, height, width = shape\n        xy = th.randint(0, 2, size=[batch, channel, height, width, 2])\n        grid = th.arange(height * width).reshape(height, width)\n        indices = grid * 2 + (th.arange(height) * width * 2)[:, None] + xy[..., 0] + width * 2 * xy[..., 1]\n        return indices.cuda()\n\n    def encode(self, inputs):\n        down1, indices_1, unpool_shape1 = self.down1(inputs)\n        down2, indices_2, unpool_shape2 = self.down2(down1)\n        down3, indices_3, unpool_shape3 = self.down3(down2)\n        down4, indices_4, unpool_shape4 = self.down4(down3)\n        down5, indices_5, unpool_shape5 = self.down5(down4)\n        return down5\n\n    def decode(self, inp):\n        batch, _, height, width = inp.shape\n        up5 = self.up5(inp, self.random_indices([batch, 512, height, width]), [batch, 512, height * 2, width * 2])\n        up4 = self.up4(\n            up5, self.random_indices([batch, 512, height * 2, width * 2]), [batch, 512, height * 4, width * 4]\n        )\n        up3 = self.up3(\n            up4, self.random_indices([batch, 256, height * 4, width * 4]), [batch, 256, height * 8, width * 8]\n        )\n        up2 = self.up2(\n            up3, self.random_indices([batch, 128, height * 8, width * 8]), [batch, 128, height * 16, width * 16]\n        )\n        up1 = self.up1(\n            up2, self.random_indices([batch, 64, height * 16, width * 16]), [batch, 64, height * 32, width * 32]\n        )\n        return up1\n\n    def forward(self, inputs):\n        down1, indices_1, unpool_shape1 = self.down1(inputs)\n        down2, indices_2, unpool_shape2 = self.down2(down1)\n        down3, indices_3, unpool_shape3 = self.down3(down2)\n        down4, indices_4, unpool_shape4 = self.down4(down3)\n        down5, indices_5, unpool_shape5 = self.down5(down4)\n\n        up5 = self.up5(down5, indices_5.shape, unpool_shape5)\n        up4 = self.up4(up5, indices_4.shape, unpool_shape4)\n        up3 = self.up3(up4, indices_3.shape, unpool_shape3)\n        up2 = self.up2(up3, indices_2.shape, unpool_shape2)\n        up1 = self.up1(up2, indices_1.shape, unpool_shape1)\n\n        return up1\n\n    def init_vgg16_params(self, vgg16):\n        blocks = [self.down1, self.down2, self.down3, self.down4, self.down5]\n\n        ranges = [[0, 4], [5, 9], [10, 16], [17, 23], [24, 29]]\n        features = list(vgg16.features.children())\n\n        vgg_layers = []\n        for _layer in features:\n            if isinstance(_layer, th.nn.Conv2d):\n                vgg_layers.append(_layer)\n\n        merged_layers = []\n        for idx, conv_block in enumerate(blocks):\n            if idx < 2:\n                units = [conv_block.conv1.cbr_unit, conv_block.conv2.cbr_unit]\n            else:\n                units = [\n                    conv_block.conv1.cbr_unit,\n                    conv_block.conv2.cbr_unit,\n                    conv_block.conv3.cbr_unit,\n                ]\n            for _unit in units:\n                for _layer in _unit:\n                    if isinstance(_layer, th.nn.Conv2d):\n                        merged_layers.append(_layer)\n\n        assert len(vgg_layers) == len(merged_layers)\n\n        for l1, l2 in zip(vgg_layers, merged_layers):\n            if isinstance(l1, th.nn.Conv2d) and isinstance(l2, th.nn.Conv2d):\n                assert l1.weight.size() == l2.weight.size()\n                assert l1.bias.size() == l2.bias.size()\n                l2.weight.data = l1.weight.data\n                l2.bias.data = l1.bias.data\n\n\nclass ConvSegNet(th.nn.Module):\n    \"\"\"\n    Adapted from https://github.com/foamliu/Autoencoder\n    See LICENSE_AUTOENCODER\n    \"\"\"\n\n    def __init__(self, in_channels=3):\n        super(ConvSegNet, self).__init__()\n\n        self.encoder = th.nn.Sequential(\n            conv2DBatchNormRelu(in_channels, 64, 3, 1, 1),\n            conv2DBatchNormRelu(64, 64, 3, 1, 1),\n            th.nn.MaxPool2d(2, 2),\n            conv2DBatchNormRelu(64, 128, 3, 1, 1),\n            conv2DBatchNormRelu(128, 128, 3, 1, 1),\n            th.nn.MaxPool2d(2, 2),\n            conv2DBatchNormRelu(128, 256, 3, 1, 1),\n            conv2DBatchNormRelu(256, 256, 3, 1, 1),\n            conv2DBatchNormRelu(256, 256, 3, 1, 1),\n            th.nn.MaxPool2d(2, 2),\n            conv2DBatchNormRelu(256, 512, 3, 1, 1),\n            conv2DBatchNormRelu(512, 512, 3, 1, 1),\n            conv2DBatchNormRelu(512, 512, 3, 1, 1),\n            th.nn.MaxPool2d(2, 2),\n            conv2DBatchNormRelu(512, 512, 3, 1, 1),\n            conv2DBatchNormRelu(512, 512, 3, 1, 1),\n            conv2DBatchNormRelu(512, 512, 3, 1, 1),\n            th.nn.MaxPool2d(2, 2),\n            th.nn.Tanh(),\n        )\n\n        self.decoder = th.nn.Sequential(\n            th.nn.Upsample(scale_factor=2, mode=\"bilinear\", align_corners=False),\n            conv2DBatchNormRelu(512, 512, 3, 1, 1),\n            conv2DBatchNormRelu(512, 512, 3, 1, 1),\n            conv2DBatchNormRelu(512, 512, 3, 1, 1),\n            th.nn.Upsample(scale_factor=2, mode=\"bilinear\", align_corners=False),\n            conv2DBatchNormRelu(512, 512, 3, 1, 1),\n            conv2DBatchNormRelu(512, 512, 3, 1, 1),\n            conv2DBatchNormRelu(512, 256, 3, 1, 1),\n            th.nn.Upsample(scale_factor=2, mode=\"bilinear\", align_corners=False),\n            conv2DBatchNormRelu(256, 256, 3, 1, 1),\n            conv2DBatchNormRelu(256, 256, 3, 1, 1),\n            conv2DBatchNormRelu(256, 128, 3, 1, 1),\n            th.nn.Upsample(scale_factor=2, mode=\"bilinear\", align_corners=False),\n            conv2DBatchNormRelu(128, 128, 3, 1, 1),\n            conv2DBatchNormRelu(128, 64, 3, 1, 1),\n            th.nn.Upsample(scale_factor=2, mode=\"bilinear\", align_corners=False),\n            conv2DBatchNormRelu(64, 64, 3, 1, 1),\n            conv2DBatchNormRelu(64, in_channels, 3, 1, 1),\n        )\n\n    def encode(self, inputs):\n        return self.encoder(inputs)\n\n    def decode(self, inputs):\n        return self.decoder(inputs)\n\n    def forward(self, inputs):\n        z = self.encode(inputs)\n        # print(z.min(), z.mean(), z.max(), z.shape)\n        return self.decode(z)\n\n\nclass VariationalConvSegNet(th.nn.Module):\n    \"\"\"\n    Adapted from https://github.com/foamliu/Autoencoder\n    See LICENSE_AUTOENCODER\n    \"\"\"\n\n    def __init__(self, in_channels=3):\n        super(VariationalConvSegNet, self).__init__()\n\n        self.encoder = th.nn.Sequential(\n            conv2DBatchNormRelu(in_channels, 64, 3, 1, 1),\n            conv2DBatchNormRelu(64, 64, 3, 1, 1),\n            th.nn.MaxPool2d(2, 2),\n            conv2DBatchNormRelu(64, 128, 3, 1, 1),\n            conv2DBatchNormRelu(128, 128, 3, 1, 1),\n            th.nn.MaxPool2d(2, 2),\n            conv2DBatchNormRelu(128, 256, 3, 1, 1),\n            conv2DBatchNormRelu(256, 256, 3, 1, 1),\n            conv2DBatchNormRelu(256, 256, 3, 1, 1),\n            th.nn.MaxPool2d(2, 2),\n            conv2DBatchNormRelu(256, 512, 3, 1, 1),\n            conv2DBatchNormRelu(512, 512, 3, 1, 1),\n            conv2DBatchNormRelu(512, 512, 3, 1, 1),\n            th.nn.MaxPool2d(2, 2),\n            conv2DBatchNormRelu(512, 512, 3, 1, 1),\n            conv2DBatchNormRelu(512, 512, 3, 1, 1),\n            conv2DBatchNormRelu(512, 512, 3, 1, 1),\n            th.nn.MaxPool2d(2, 2),\n            th.nn.Tanh(),\n            Flatten(),\n        )\n\n        self.fc_mu = th.nn.Linear(512 * 4 * 4, 512 * 4 * 4)\n        self.fc_var = th.nn.Linear(512 * 4 * 4, 512 * 4 * 4)\n\n        self.decoder = th.nn.Sequential(\n            UnFlatten(512, 4),\n            th.nn.Upsample(scale_factor=2, mode=\"bilinear\", align_corners=False),\n            conv2DBatchNormRelu(512, 512, 3, 1, 1),\n            conv2DBatchNormRelu(512, 512, 3, 1, 1),\n            conv2DBatchNormRelu(512, 512, 3, 1, 1),\n            th.nn.Upsample(scale_factor=2, mode=\"bilinear\", align_corners=False),\n            conv2DBatchNormRelu(512, 512, 3, 1, 1),\n            conv2DBatchNormRelu(512, 512, 3, 1, 1),\n            conv2DBatchNormRelu(512, 256, 3, 1, 1),\n            th.nn.Upsample(scale_factor=2, mode=\"bilinear\", align_corners=False),\n            conv2DBatchNormRelu(256, 256, 3, 1, 1),\n            conv2DBatchNormRelu(256, 256, 3, 1, 1),\n            conv2DBatchNormRelu(256, 128, 3, 1, 1),\n            th.nn.Upsample(scale_factor=2, mode=\"bilinear\", align_corners=False),\n            conv2DBatchNormRelu(128, 128, 3, 1, 1),\n            conv2DBatchNormRelu(128, 64, 3, 1, 1),\n            th.nn.Upsample(scale_factor=2, mode=\"bilinear\", align_corners=False),\n            conv2DBatchNormRelu(64, 64, 3, 1, 1),\n            conv2DBatchNormRelu(64, in_channels, 3, 1, 1),\n            th.nn.Tanh(),\n        )\n\n    def reparameterize(self, mu, log_var):\n        std = th.exp(0.5 * log_var)\n        eps = th.randn_like(std)\n        return eps * std + mu\n\n    def encode(self, inputs):\n        result = self.encoder(inputs)\n\n        mu = self.fc_mu(result)\n        log_var = self.fc_var(result)\n\n        return mu, log_var\n\n    def decode(self, inputs):\n        return self.decoder(inputs)\n\n    def forward(self, inputs):\n        mu, log_var = self.encode(inputs)\n        z = self.reparameterize(mu, log_var)\n        return self.decode(z)\n\n\ndef create_encoder_single_conv(in_chs, out_chs, kernel):\n    assert kernel % 2 == 1\n    return th.nn.Sequential(\n        th.nn.Conv2d(in_chs, out_chs, kernel_size=kernel, padding=(kernel - 1) // 2),\n        th.nn.BatchNorm2d(out_chs),\n        FusedLeakyReLU(out_chs),\n    )\n\n\nclass EncoderInceptionModuleSignle(th.nn.Module):\n    def __init__(self, channels):\n        assert channels % 2 == 0\n        super().__init__()\n        # put bottle-neck layers before convolution\n        bn_ch = channels // 2\n        self.bottleneck = create_encoder_single_conv(channels, bn_ch, 1)\n        # bn -> Conv1, 3, 5\n        self.conv1 = create_encoder_single_conv(bn_ch, channels, 1)\n        self.conv3 = create_encoder_single_conv(bn_ch, channels, 3)\n        self.conv5 = create_encoder_single_conv(bn_ch, channels, 5)\n        self.conv7 = create_encoder_single_conv(bn_ch, channels, 7)\n        # pool-proj(no-bottle neck)\n        self.pool3 = th.nn.MaxPool2d(3, stride=1, padding=1)\n        self.pool5 = th.nn.MaxPool2d(5, stride=1, padding=2)\n\n    def forward(self, x):\n        # Original inception is concatenation, but use simple addition instead\n        bn = self.bottleneck(x)\n        out = self.conv1(bn) + self.conv3(bn) + self.conv5(bn) + self.conv7(bn) + self.pool3(x) + self.pool5(x)\n        return out\n\n\nclass EncoderModule(th.nn.Module):\n    def __init__(self, chs, repeat_num, use_inception):\n        super().__init__()\n        if use_inception:\n            layers = [EncoderInceptionModuleSignle(chs) for i in range(repeat_num)]\n        else:\n            layers = [create_encoder_single_conv(chs, chs, 3) for i in range(repeat_num)]\n        self.convs = th.nn.Sequential(*layers)\n\n    def forward(self, x):\n        return self.convs(x)\n\n\nclass Encoder(th.nn.Module):\n    def __init__(self, use_inception, repeat_per_module):\n        super().__init__()\n        # stages\n        self.upch1 = th.nn.Conv2d(3, 32, kernel_size=3)\n        self.stage1 = EncoderModule(32, repeat_per_module, use_inception)\n        self.upch2 = self._create_downsampling_module(32, 2)\n        self.stage2 = EncoderModule(64, repeat_per_module, use_inception)\n        self.upch3 = self._create_downsampling_module(64, 2)\n        self.stage3 = EncoderModule(128, repeat_per_module, use_inception)\n        self.upch4 = self._create_downsampling_module(128, 2)\n        self.stage4 = EncoderModule(256, repeat_per_module, use_inception)\n\n    def _create_downsampling_module(self, input_channels, pooling_kenel):\n        return th.nn.Sequential(\n            th.nn.AvgPool2d(pooling_kenel),\n            th.nn.Conv2d(input_channels, input_channels * 2, kernel_size=1),\n            th.nn.BatchNorm2d(input_channels * 2),\n            FusedLeakyReLU(input_channels * 2),\n        )\n\n    def forward(self, x):\n        # print(x.shape)\n        out = self.stage1(self.upch1(x))\n        # print(out.shape)\n        out = self.stage2(self.upch2(out))\n        # print(out.shape)\n        out = self.stage3(self.upch3(out))\n        # print(out.shape)\n        out = self.stage4(self.upch4(out))\n        # print(out.shape)\n        out = F.avg_pool2d(out, 8)  # Global Average pooling\n        # print(out.shape)\n        return out.view(-1, 256)\n\n\n## Decoder\ndef create_decoder_single_conv(in_chs, out_chs, kernel):\n    assert kernel % 2 == 1\n    return th.nn.Sequential(\n        th.nn.ConvTranspose2d(in_chs, out_chs, kernel_size=kernel, padding=(kernel - 1) // 2),\n        th.nn.BatchNorm2d(out_chs),\n        FusedLeakyReLU(out_chs),\n    )\n\n\nclass DecoderInceptionModuleSingle(th.nn.Module):\n    def __init__(self, channels):\n        assert channels % 2 == 0\n        super().__init__()\n        # put bottle-neck layers before convolution\n        bn_ch = channels // 4\n        self.bottleneck = create_decoder_single_conv(channels, bn_ch, 1)\n        # bn -> Conv1, 3, 5\n        self.conv1 = create_decoder_single_conv(bn_ch, channels, 1)\n        self.conv3 = create_decoder_single_conv(bn_ch, channels, 3)\n        self.conv5 = create_decoder_single_conv(bn_ch, channels, 5)\n        self.conv7 = create_decoder_single_conv(bn_ch, channels, 7)\n        # pool-proj(no-bottle neck)\n        self.pool3 = th.nn.MaxPool2d(3, stride=1, padding=1)\n        self.pool5 = th.nn.MaxPool2d(5, stride=1, padding=2)\n\n    def forward(self, x):\n        # Original inception is concatenation, but use simple addition instead\n        bn = self.bottleneck(x)\n        out = self.conv1(bn) + self.conv3(bn) + self.conv5(bn) + self.conv7(bn) + self.pool3(x) + self.pool5(x)\n        return out\n\n\nclass DecoderModule(th.nn.Module):\n    def __init__(self, chs, repeat_num, use_inception):\n        super().__init__()\n        if use_inception:\n            layers = [DecoderInceptionModuleSingle(chs) for i in range(repeat_num)]\n        else:\n            layers = [create_decoder_single_conv(chs, chs, 3) for i in range(repeat_num)]\n        self.convs = th.nn.Sequential(*layers)\n\n    def forward(self, x):\n        return self.convs(x)\n\n\nclass Decoder(th.nn.Module):\n    def __init__(self, use_inception, repeat_per_module):\n        super().__init__()\n        # stages\n        self.stage1 = DecoderModule(256, repeat_per_module, use_inception)\n        self.downch1 = self._create_upsampling_module(256, 2)\n        self.stage2 = DecoderModule(128, repeat_per_module, use_inception)\n        self.downch2 = self._create_upsampling_module(128, 2)\n        self.stage3 = DecoderModule(64, repeat_per_module, use_inception)\n        self.downch3 = self._create_upsampling_module(64, 2)\n        self.stage4 = DecoderModule(32, repeat_per_module, use_inception)\n        self.downch4 = self._create_upsampling_module(32, 2)\n        self.last = th.nn.ConvTranspose2d(16, 3, kernel_size=1)\n\n    def _create_upsampling_module(self, input_channels, pooling_kenel):\n        return th.nn.Sequential(\n            th.nn.ConvTranspose2d(input_channels, input_channels // 2, kernel_size=pooling_kenel, stride=pooling_kenel),\n            th.nn.BatchNorm2d(input_channels // 2),\n            FusedLeakyReLU(input_channels // 2),\n        )\n\n    def forward(self, x):\n        out = F.upsample(x.view(-1, 256, 1, 1), scale_factor=8)\n        out = self.downch1(self.stage1(out))\n        out = self.downch2(self.stage2(out))\n        out = self.downch3(self.stage3(out))\n        out = self.downch4(self.stage4(out))\n        return th.sigmoid(self.last(out))\n\n\n## VAE\nclass InceptionVAE(th.nn.Module):\n    \"\"\"\n    Adapted from https://github.com/koshian2/inception-vae\n    \"\"\"\n\n    def __init__(self, latent_dim=512, repeat_per_block=1, use_inception=True):\n        super(InceptionVAE, self).__init__()\n\n        # # latent features\n        self.n_latent_features = latent_dim\n\n        # Encoder\n        self.encoder = Encoder(use_inception, repeat_per_block)\n        # Middle\n        self.fc_mu = th.nn.Linear(256, self.n_latent_features)\n        self.fc_logvar = th.nn.Linear(256, self.n_latent_features)\n        self.fc_rep = th.nn.Linear(self.n_latent_features, 256)\n        # Decoder\n        self.decoder = Decoder(use_inception, repeat_per_block)\n\n    def _reparameterize(self, mu, logvar):\n        std = logvar.mul(0.5).exp_()\n        esp = th.randn(*mu.size()).cuda()\n        z = mu + std * esp\n        return z\n\n    def _bottleneck(self, h):\n        mu, logvar = self.fc_mu(h), self.fc_logvar(h)\n        z = self._reparameterize(mu, logvar)\n        return z, mu, logvar\n\n    def sampling(self):\n        # assume latent features space ~ N(0, 1)\n        z = th.randn(24, self.n_latent_features).cuda()\n        z = self.fc_rep(z)\n        # decode\n        return self.decoder(z)\n\n    def forward(self, x):\n        # Encoder\n        h = self.encoder(x)\n        # Bottle-neck\n        z, mu, logvar = self._bottleneck(h)\n        # decoder\n        z = self.fc_rep(z)\n        d = self.decoder(z)\n        return d, mu, logvar\n"
  },
  {
    "path": "models/stylegan1.py",
    "content": "# from https://github.com/lernapparat/lernapparat/blob/master/style_gan/pyth_style_gan.ipynb\n\nimport gc\nfrom collections import OrderedDict\n\nimport numpy as np\nimport torch as th\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n\nclass MyLinear(nn.Module):\n    \"\"\"Linear layer with equalized learning rate and custom learning rate multiplier.\"\"\"\n\n    def __init__(\n        self, input_size, output_size, gain=2 ** (0.5), use_wscale=False, lrmul=1, bias=True,\n    ):\n        super().__init__()\n        he_std = gain * input_size ** (-0.5)  # He init\n        # Equalized learning rate and custom learning rate multiplier.\n        if use_wscale:\n            init_std = 1.0 / lrmul\n            self.w_mul = he_std * lrmul\n        else:\n            init_std = he_std / lrmul\n            self.w_mul = lrmul\n        self.weight = th.nn.Parameter(th.randn(output_size, input_size) * init_std)\n        if bias:\n            self.bias = th.nn.Parameter(th.zeros(output_size))\n            self.b_mul = lrmul\n        else:\n            self.bias = None\n\n    def forward(self, x):\n        bias = self.bias\n        if bias is not None:\n            bias = bias * self.b_mul\n        return F.linear(x, self.weight * self.w_mul, bias)\n\n\nclass MyConv2d(nn.Module):\n    \"\"\"Conv layer with equalized learning rate and custom learning rate multiplier.\"\"\"\n\n    def __init__(\n        self,\n        input_channels,\n        output_channels,\n        kernel_size,\n        gain=2 ** (0.5),\n        use_wscale=False,\n        lrmul=1,\n        bias=True,\n        intermediate=None,\n        upscale=False,\n    ):\n        super().__init__()\n        if upscale:\n            self.upscale = Upscale2d()\n        else:\n            self.upscale = None\n        he_std = gain * (input_channels * kernel_size ** 2) ** (-0.5)  # He init\n        self.kernel_size = kernel_size\n        if use_wscale:\n            init_std = 1.0 / lrmul\n            self.w_mul = he_std * lrmul\n        else:\n            init_std = he_std / lrmul\n            self.w_mul = lrmul\n        self.weight = th.nn.Parameter(th.randn(output_channels, input_channels, kernel_size, kernel_size) * init_std)\n        if bias:\n            self.bias = th.nn.Parameter(th.zeros(output_channels))\n            self.b_mul = lrmul\n        else:\n            self.bias = None\n        self.intermediate = intermediate\n\n    def forward(self, x):\n        bias = self.bias\n        if bias is not None:\n            bias = bias * self.b_mul\n\n        have_convolution = False\n        if self.upscale is not None and min(x.shape[2:]) * 2 >= 128:\n            # this is the fused upscale + conv from StyleGAN, sadly this seems incompatible with the non-fused way\n            # this really needs to be cleaned up and go into the conv...\n            w = self.weight * self.w_mul\n            w = w.permute(1, 0, 2, 3)\n            # probably applying a conv on w would be more efficient. also this quadruples the weight (average)?!\n            w = F.pad(w, (1, 1, 1, 1))\n            w = w[:, :, 1:, 1:] + w[:, :, :-1, 1:] + w[:, :, 1:, :-1] + w[:, :, :-1, :-1]\n            x = F.conv_transpose2d(x, w, stride=2, padding=(w.size(-1) - 1) // 2)\n            have_convolution = True\n        elif self.upscale is not None:\n            x = self.upscale(x)\n\n        if not have_convolution and self.intermediate is None:\n            return F.conv2d(x, self.weight * self.w_mul, bias, padding=self.kernel_size // 2)\n        elif not have_convolution:\n            x = F.conv2d(x, self.weight * self.w_mul, None, padding=self.kernel_size // 2)\n\n        if self.intermediate is not None:\n            x = self.intermediate(x)\n        if bias is not None:\n            x = x + bias.view(1, -1, 1, 1)\n        return x\n\n\nclass NoiseLayer(nn.Module):\n    \"\"\"adds noise. noise is per pixel (constant over channels) with per-channel weight\"\"\"\n\n    def __init__(self, channels):\n        super().__init__()\n        self.weight = nn.Parameter(th.zeros(channels))\n        self.noise = None\n\n    def forward(self, x):\n        if self.noise is None:\n            noise = th.randn(x.size(0), 1, x.size(2), x.size(3), device=x.device, dtype=x.dtype)\n        else:\n            noise = self.noise.to(x.device)\n        # print(noise.shape, noise.min(), noise.mean(), noise.max())\n        x = x + self.weight.view(1, -1, 1, 1) * noise\n        return x\n\n\nclass StyleMod(nn.Module):\n    def __init__(self, latent_size, channels, use_wscale):\n        super(StyleMod, self).__init__()\n        self.lin = MyLinear(latent_size, channels * 2, gain=1.0, use_wscale=use_wscale)\n\n    def forward(self, x, latent):\n        style = self.lin(latent)  # style => [batch_size, n_channels*2]\n        shape = [-1, 2, x.size(1)] + (x.dim() - 2) * [1]\n        style = style.view(shape)  # [batch_size, 2, n_channels, ...]\n        x = x * (style[:, 0] + 1.0) + style[:, 1]\n        return x\n\n\nclass PixelNormLayer(nn.Module):\n    def __init__(self, epsilon=1e-8):\n        super().__init__()\n        self.epsilon = epsilon\n\n    def forward(self, x):\n        return x * th.rsqrt(th.mean(x ** 2, dim=1, keepdim=True) + self.epsilon)\n\n\nclass BlurLayer(nn.Module):\n    def __init__(self, kernel=[1, 2, 1], normalize=True, flip=False, stride=1):\n        super(BlurLayer, self).__init__()\n        kernel = [1, 2, 1]\n        kernel = th.tensor(kernel, dtype=th.float32)\n        kernel = kernel[:, None] * kernel[None, :]\n        kernel = kernel[None, None]\n        if normalize:\n            kernel = kernel / kernel.sum()\n        if flip:\n            kernel = kernel[:, :, ::-1, ::-1]\n        self.register_buffer(\"kernel\", kernel)\n        self.stride = stride\n\n    def forward(self, x):\n        # expand kernel channels\n        kernel = self.kernel.expand(x.size(1), -1, -1, -1)\n        x = F.conv2d(x, kernel, stride=self.stride, padding=int((self.kernel.size(2) - 1) / 2), groups=x.size(1),)\n        return x\n\n\ndef upscale2d(x, factor=2, gain=1):\n    assert x.dim() == 4\n    if gain != 1:\n        x = x * gain\n    if factor != 1:\n        shape = x.shape\n        x = x.view(shape[0], shape[1], shape[2], 1, shape[3], 1).expand(-1, -1, -1, factor, -1, factor)\n        x = x.contiguous().view(shape[0], shape[1], factor * shape[2], factor * shape[3])\n    return x\n\n\nclass Upscale2d(nn.Module):\n    def __init__(self, factor=2, gain=1):\n        super().__init__()\n        assert isinstance(factor, int) and factor >= 1\n        self.gain = gain\n        self.factor = factor\n\n    def forward(self, x):\n        return upscale2d(x, factor=self.factor, gain=self.gain)\n\n\nclass G_mapping(nn.Sequential):\n    def __init__(self, nonlinearity=\"lrelu\", use_wscale=True):\n        act, gain = {\"relu\": (th.relu, np.sqrt(2)), \"lrelu\": (nn.LeakyReLU(negative_slope=0.2), np.sqrt(2))}[\n            nonlinearity\n        ]\n        layers = [\n            (\"pixel_norm\", PixelNormLayer()),\n            (\"dense0\", MyLinear(512, 512, gain=gain, lrmul=0.01, use_wscale=use_wscale),),\n            (\"dense0_act\", act),\n            (\"dense1\", MyLinear(512, 512, gain=gain, lrmul=0.01, use_wscale=use_wscale),),\n            (\"dense1_act\", act),\n            (\"dense2\", MyLinear(512, 512, gain=gain, lrmul=0.01, use_wscale=use_wscale),),\n            (\"dense2_act\", act),\n            (\"dense3\", MyLinear(512, 512, gain=gain, lrmul=0.01, use_wscale=use_wscale),),\n            (\"dense3_act\", act),\n            (\"dense4\", MyLinear(512, 512, gain=gain, lrmul=0.01, use_wscale=use_wscale),),\n            (\"dense4_act\", act),\n            (\"dense5\", MyLinear(512, 512, gain=gain, lrmul=0.01, use_wscale=use_wscale),),\n            (\"dense5_act\", act),\n            (\"dense6\", MyLinear(512, 512, gain=gain, lrmul=0.01, use_wscale=use_wscale),),\n            (\"dense6_act\", act),\n            (\"dense7\", MyLinear(512, 512, gain=gain, lrmul=0.01, use_wscale=use_wscale),),\n            (\"dense7_act\", act),\n        ]\n        super().__init__(OrderedDict(layers))\n\n    def forward(self, x):\n        x = super().forward(x)\n        # Broadcast\n        x = x.unsqueeze(1).expand(-1, 18, -1)\n        return x\n\n\nclass Truncation(nn.Module):\n    def __init__(self, avg_latent, max_layer=8, threshold=0.7):\n        super().__init__()\n        self.max_layer = max_layer\n        self.threshold = threshold\n        self.register_buffer(\"avg_latent\", avg_latent)\n\n    def forward(self, x):\n        assert x.dim() == 3\n        interp = th.lerp(self.avg_latent, x, self.threshold)\n        do_trunc = (th.arange(x.size(1)) < self.max_layer).view(1, -1, 1)\n        return th.where(do_trunc, interp, x)\n\n\nclass LayerEpilogue(nn.Module):\n    \"\"\"Things to do at the end of each layer.\"\"\"\n\n    def __init__(\n        self,\n        channels,\n        dlatent_size,\n        use_wscale,\n        use_noise,\n        use_pixel_norm,\n        use_instance_norm,\n        use_styles,\n        activation_layer,\n    ):\n        super().__init__()\n        layers = []\n        if use_noise:\n            layers.append((\"noise\", NoiseLayer(channels)))\n        layers.append((\"activation\", activation_layer))\n        if use_pixel_norm:\n            layers.append((\"pixel_norm\", PixelNormLayer()))\n        if use_instance_norm:\n            layers.append((\"instance_norm\", nn.InstanceNorm2d(channels)))\n        self.top_epi = nn.Sequential(OrderedDict(layers))\n        if use_styles:\n            self.style_mod = StyleMod(dlatent_size, channels, use_wscale=use_wscale)\n        else:\n            self.style_mod = None\n\n        # if use_noise:\n        #     # layers.append((\"noise\", NoiseLayer(channels)))\n        #     self.noise = NoiseLayer(channels)\n        # else:\n        #     self.noise = None\n\n        # self.activation = activation_layer\n\n        # if use_pixel_norm:\n        #     self.pixel_norm = PixelNormLayer()\n        # else:\n        #     self.pixel_norm = None\n\n        # if use_instance_norm:\n        #     self.instance_norm = nn.InstanceNorm2d(channels)\n        # else:\n        #     self.instance_norm = None\n\n        # if use_styles:\n        #     self.style_mod = StyleMod(dlatent_size, channels, use_wscale=use_wscale)\n        # else:\n        #     self.style_mod = None\n\n    def forward(self, x, dlatents_in_slice, noise):\n        # if self.noise is not None:\n        #     x = self.noise(x, noise)\n\n        # x = self.activation(x)\n\n        # if self.pixel_norm is not None:\n        #     x = self.pixel_norm(x)\n\n        # if self.instance_norm is not None:\n        #     x = self.instance_norm(x)\n\n        if noise is not None:\n            self.top_epi.noise.noise = noise\n\n        x = self.top_epi(x)\n        if self.style_mod is not None:\n            x = self.style_mod(x, dlatents_in_slice)\n        else:\n            assert dlatents_in_slice is None\n\n        if noise is not None:\n            del self.top_epi.noise.noise\n            gc.collect()\n\n        return x\n\n\nclass InputBlock(nn.Module):\n    def __init__(\n        self,\n        nf,\n        dlatent_size,\n        const_input_layer,\n        gain,\n        use_wscale,\n        use_noise,\n        use_pixel_norm,\n        use_instance_norm,\n        use_styles,\n        activation_layer,\n    ):\n        super().__init__()\n        self.const_input_layer = const_input_layer\n        self.nf = nf\n        if self.const_input_layer:\n            # called 'const' in tf\n            self.const = nn.Parameter(th.ones(1, nf, 4, 4))\n            self.bias = nn.Parameter(th.ones(nf))\n        else:\n            self.dense = MyLinear(\n                dlatent_size, nf * 16, gain=gain / 4, use_wscale=use_wscale\n            )  # tweak gain to match the official implementation of Progressing GAN\n        self.epi1 = LayerEpilogue(\n            nf, dlatent_size, use_wscale, use_noise, use_pixel_norm, use_instance_norm, use_styles, activation_layer,\n        )\n        self.conv = MyConv2d(nf, nf, 3, gain=gain, use_wscale=use_wscale)\n        self.epi2 = LayerEpilogue(\n            nf, dlatent_size, use_wscale, use_noise, use_pixel_norm, use_instance_norm, use_styles, activation_layer,\n        )\n\n    def forward(self, dlatents_in_range, noise):\n        batch_size = dlatents_in_range.size(0)\n        if self.const_input_layer:\n            x = self.const.expand(batch_size, -1, -1, -1)\n            x = x + self.bias.view(1, -1, 1, 1)\n        else:\n            x = self.dense(dlatents_in_range[:, 0]).view(batch_size, self.nf, 4, 4)\n        x = self.epi1(x, dlatents_in_range[:, 0], noise=noise)\n        x = self.conv(x)\n        x = self.epi2(x, dlatents_in_range[:, 1], noise=noise)\n        return x\n\n\nclass GSynthesisBlock(nn.Module):\n    def __init__(\n        self,\n        in_channels,\n        out_channels,\n        blur_filter,\n        dlatent_size,\n        gain,\n        use_wscale,\n        use_noise,\n        use_pixel_norm,\n        use_instance_norm,\n        use_styles,\n        activation_layer,\n    ):\n        # 2**res x 2**res # res = 3..resolution_log2\n        super().__init__()\n        if blur_filter:\n            blur = BlurLayer(blur_filter)\n        else:\n            blur = None\n        self.conv0_up = MyConv2d(\n            in_channels, out_channels, kernel_size=3, gain=gain, use_wscale=use_wscale, intermediate=blur, upscale=True,\n        )\n        self.epi1 = LayerEpilogue(\n            out_channels,\n            dlatent_size,\n            use_wscale,\n            use_noise,\n            use_pixel_norm,\n            use_instance_norm,\n            use_styles,\n            activation_layer,\n        )\n        self.conv1 = MyConv2d(out_channels, out_channels, kernel_size=3, gain=gain, use_wscale=use_wscale)\n        self.epi2 = LayerEpilogue(\n            out_channels,\n            dlatent_size,\n            use_wscale,\n            use_noise,\n            use_pixel_norm,\n            use_instance_norm,\n            use_styles,\n            activation_layer,\n        )\n\n    def forward(self, x, dlatents_in_range, noise):\n        x = self.conv0_up(x)\n        x = self.epi1(x, dlatents_in_range[:, 0], noise=noise)\n        x = self.conv1(x)\n        x = self.epi2(x, dlatents_in_range[:, 1], noise=noise)\n        return x\n\n\nclass G_synthesis(nn.Module):\n    def __init__(\n        self,\n        dlatent_size=512,  # Disentangled latent (W) dimensionality.\n        num_channels=3,  # Number of output color channels.\n        resolution=1024,  # Output resolution.\n        fmap_base=8192,  # Overall multiplier for the number of feature maps.\n        fmap_decay=1.0,  # log2 feature map reduction when doubling the resolution.\n        fmap_max=512,  # Maximum number of feature maps in any layer.\n        use_styles=True,  # Enable style inputs?\n        const_input_layer=True,  # First layer is a learned constant?\n        use_noise=True,  # Enable noise inputs?\n        randomize_noise=False,  # True = randomize noise inputs every time (non-deterministic) or from variables passed,\n        nonlinearity=\"lrelu\",  # Activation function: 'relu', 'lrelu'\n        use_wscale=True,  # Enable equalized learning rate?\n        use_pixel_norm=False,  # Enable pixelwise feature vector normalization?\n        use_instance_norm=True,  # Enable instance normalization?\n        dtype=th.float32,  # Data type to use for activations and outputs.\n        blur_filter=[1, 2, 1],  # Low-pass filter to apply when resampling activations. None = no filtering.\n    ):\n\n        super().__init__()\n\n        def nf(stage):\n            return min(int(fmap_base / (2.0 ** (stage * fmap_decay))), fmap_max)\n\n        self.dlatent_size = dlatent_size\n        resolution_log2 = int(np.log2(resolution))\n        assert resolution == 2 ** resolution_log2 and resolution >= 4\n\n        act, gain = {\"relu\": (th.relu, np.sqrt(2)), \"lrelu\": (nn.LeakyReLU(negative_slope=0.2), np.sqrt(2)),}[\n            nonlinearity\n        ]\n        blocks = []\n        for res in range(2, resolution_log2 + 1):\n            channels = nf(res - 1)\n            name = \"{s}x{s}\".format(s=2 ** res)\n            if res == 2:\n                blocks.append(\n                    (\n                        name,\n                        InputBlock(\n                            channels,\n                            dlatent_size,\n                            const_input_layer,\n                            gain,\n                            use_wscale,\n                            use_noise,\n                            use_pixel_norm,\n                            use_instance_norm,\n                            use_styles,\n                            act,\n                        ),\n                    )\n                )\n\n            else:\n                blocks.append(\n                    (\n                        name,\n                        GSynthesisBlock(\n                            last_channels,\n                            channels,\n                            blur_filter,\n                            dlatent_size,\n                            gain,\n                            use_wscale,\n                            use_noise,\n                            use_pixel_norm,\n                            use_instance_norm,\n                            use_styles,\n                            act,\n                        ),\n                    )\n                )\n            last_channels = channels\n        self.torgb = MyConv2d(channels, num_channels, 1, gain=1, use_wscale=use_wscale)\n        self.blocks = nn.ModuleDict(OrderedDict(blocks))\n\n    def forward(self, dlatents_in, noise):\n        # Input: Disentangled latents (W) [minibatch, num_layers, dlatent_size].\n        # lod_in = tf.cast(tf.get_variable('lod', initializer=np.float32(0), trainable=False), dtype)\n        for i, m in enumerate(self.blocks.values()):\n            if i == 0:\n                x = m(dlatents_in[:, 2 * i : 2 * i + 2], noise=noise)\n            else:\n                x = m(x, dlatents_in[:, 2 * i : 2 * i + 2], noise=noise)\n        rgb = self.torgb(x)\n        return rgb\n\n\nclass G_style(nn.Sequential):\n    def __init__(self, output_size=1920, checkpoint=None):\n\n        # TODO FIX THIS MONSTROSITY\n        super().__init__()\n        self.g_mapping = G_mapping()\n        try:\n            self.g_synthesis = G_synthesis(resolution=1024)\n            if checkpoint is not None:\n                self.load_state_dict(th.load(checkpoint), strict=False)\n            network_resolution = 1024\n        except:\n            print(\"Trying 512px generator resolution...\")\n            try:\n                self.g_synthesis = G_synthesis(resolution=512)\n                if checkpoint is not None:\n                    self.load_state_dict(th.load(checkpoint), strict=False)\n                network_resolution = 512\n            except:\n                print(\"Trying 256px generator resolution...\")\n                try:\n                    self.g_synthesis = G_synthesis(resolution=256)\n                    if checkpoint is not None:\n                        self.load_state_dict(th.load(checkpoint), strict=False)\n                    network_resolution = 256\n                except:\n                    print(\"Trying 128px generator resolution...\")\n                    try:\n                        self.g_synthesis = G_synthesis(resolution=128)\n                        if checkpoint is not None:\n                            self.load_state_dict(th.load(checkpoint), strict=False)\n                        network_resolution = 128\n                    except:\n                        print(\"ERROR: Network too small or state_dict mismatch\")\n                        exit()\n\n        const = getattr(self.g_synthesis.blocks, \"4x4\").const\n        if network_resolution != 1024:\n            means = th.zeros(size=(1, 512, int(4 * 1024 / network_resolution), int(4 * 1024 / network_resolution)))\n            const = th.normal(mean=means, std=th.ones_like(means) * const.std(),)\n\n        _, _, ch, cw = const.shape\n        if output_size == 1920:\n            layer0 = th.cat(\n                [\n                    const[:, :, :, [0]],\n                    const[:, :, :, [0]],\n                    # const[:, :, :, : cw // 2 + 1][:, :, :, list(range(cw // 2, 0, -1))],\n                    const,\n                    # const[:, :, :, cw // 2 :],\n                    const[:, :, :, [-1]],\n                    const[:, :, :, [-1]],\n                ],\n                axis=3,\n            )\n        elif output_size == 512:\n            layer0 = const[:, :, ch // 4 : 3 * ch // 4, cw // 4 : 3 * cw // 4]\n        else:\n            layer0 = const\n        getattr(self.g_synthesis.blocks, \"4x4\").const = th.nn.Parameter(layer0 + th.normal(0, const.std() / 2.0))\n        _, _, height, width = getattr(self.g_synthesis.blocks, \"4x4\").const.shape\n\n        for i in range(len(list(self.g_synthesis.blocks.named_parameters())) // 10):\n            self.register_buffer(f\"noise_{i}\", th.randn(1, 1, height * 2 ** i, width * 2 ** i))\n\n        self.truncation_latent = self.mean_latent(2 ** 14)\n\n    def mean_latent(self, n_latent):\n        latent_in = th.randn(n_latent, 512)\n        latent = self.g_mapping(latent_in).mean(0, keepdim=True)\n        return latent\n\n    def forward(\n        self,\n        styles,\n        noise=None,\n        truncation=1,\n        map_latents=False,\n        randomize_noise=False,\n        input_is_latent=True,\n        transform_dict_list=None,\n    ):\n        if map_latents:\n            return self.g_mapping(styles)\n\n        if noise is None:\n            noise = [None] * (len(list(self.g_synthesis.blocks.named_parameters())) // 10)\n        for ns, noise_scale in enumerate(noise):\n            if noise_scale is None:\n                try:\n                    noise[ns] = getattr(self, f\"noise_{ns}\")\n                except:\n                    pass\n\n        if truncation != 1:\n            interp = th.lerp(self.truncation_latent.to(styles.device), styles, truncation)\n            do_trunc = (th.arange(styles.size(1)) < 8).view(1, -1, 1).to(styles.device)\n            styles = th.where(do_trunc, interp, styles)\n\n        # Input: Disentangled latents (W) [minibatch, num_layers, dlatent_size].\n        # print(styles.shape, len(noise), len(self.g_synthesis.blocks.values()))\n        for i, block in enumerate(self.g_synthesis.blocks.values()):\n            if i == 0:\n                x = block(styles[:, 2 * i : 2 * i + 2], noise=noise[i])\n            else:\n                x = block(x, styles[:, 2 * i : 2 * i + 2], noise=noise[i])\n        img = self.g_synthesis.torgb(x)\n\n        return img, None\n"
  },
  {
    "path": "models/stylegan2.py",
    "content": "import math\nimport os\nimport random\nimport sys\n\nimport torch as th\nfrom torch import nn\nfrom torch.nn import functional as F\n\nsys.path.append(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))\n\nfrom op import FusedLeakyReLU, fused_leaky_relu, upfirdn2d\n\n\nclass PixelNorm(nn.Module):\n    def __init__(self):\n        super().__init__()\n\n    def forward(self, inputs):\n        return inputs * th.rsqrt(th.mean(inputs ** 2, dim=1, keepdim=True) + 1e-8)\n\n\ndef make_kernel(k):\n    k = th.tensor(k, dtype=th.float32)\n\n    if k.ndim == 1:\n        k = k[None, :] * k[:, None]\n\n    k /= k.sum()\n\n    return k\n\n\nclass Upsample(nn.Module):\n    def __init__(self, kernel, factor=2):\n        super().__init__()\n\n        self.factor = factor\n        kernel = make_kernel(kernel) * (factor ** 2)\n        self.register_buffer(\"kernel\", kernel)\n\n        p = kernel.shape[0] - factor\n\n        pad0 = (p + 1) // 2 + factor - 1\n        pad1 = p // 2\n\n        self.pad = (pad0, pad1)\n\n    def forward(self, inputs):\n        out = upfirdn2d(inputs, self.kernel, up=self.factor, down=1, pad=self.pad)\n\n        return out\n\n\nclass Downsample(nn.Module):\n    def __init__(self, kernel, factor=2):\n        super().__init__()\n\n        self.factor = factor\n        kernel = make_kernel(kernel)\n        self.register_buffer(\"kernel\", kernel)\n\n        p = kernel.shape[0] - factor\n\n        pad0 = (p + 1) // 2\n        pad1 = p // 2\n\n        self.pad = (pad0, pad1)\n\n    def forward(self, inputs):\n        out = upfirdn2d(inputs, self.kernel, up=1, down=self.factor, pad=self.pad)\n\n        return out\n\n\nclass Blur(nn.Module):\n    def __init__(self, kernel, pad, upsample_factor=1):\n        super().__init__()\n\n        kernel = make_kernel(kernel)\n\n        if upsample_factor > 1:\n            kernel = kernel * (upsample_factor ** 2)\n\n        self.register_buffer(\"kernel\", kernel)\n\n        self.pad = pad\n\n    def forward(self, inputs):\n        out = upfirdn2d(inputs, self.kernel, pad=self.pad)\n\n        return out\n\n\nclass EqualConv2d(nn.Module):\n    def __init__(self, in_channel, out_channel, kernel_size, stride=1, padding=0, bias=True):\n        super().__init__()\n\n        self.weight = nn.Parameter(th.randn(out_channel, in_channel, kernel_size, kernel_size))\n        self.scale = 1 / math.sqrt(in_channel * kernel_size ** 2)\n\n        self.stride = stride\n        self.padding = padding\n\n        if bias:\n            self.bias = nn.Parameter(th.zeros(out_channel))\n\n        else:\n            self.bias = None\n\n    def forward(self, inputs):\n        out = F.conv2d(inputs, self.weight * self.scale, bias=self.bias, stride=self.stride, padding=self.padding,)\n\n        return out\n\n    def __repr__(self):\n        return (\n            f\"{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]},\"\n            f\" {self.weight.shape[2]}, stride={self.stride}, padding={self.padding})\"\n        )\n\n\nclass EqualLinear(nn.Module):\n    def __init__(self, in_dim, out_dim, bias=True, bias_init=0, lr_mul=1, activation=None):\n        super().__init__()\n\n        self.weight = nn.Parameter(th.randn(out_dim, in_dim).div_(lr_mul))\n\n        if bias:\n            self.bias = nn.Parameter(th.zeros(out_dim).fill_(bias_init))\n\n        else:\n            self.bias = None\n\n        self.activation = activation\n\n        self.scale = (1 / math.sqrt(in_dim)) * lr_mul\n        self.lr_mul = lr_mul\n\n    def forward(self, inputs):\n        if self.activation:\n            out = F.linear(inputs, self.weight * self.scale)\n            out = fused_leaky_relu(out, self.bias * self.lr_mul)\n        else:\n            out = F.linear(inputs, self.weight * self.scale, bias=self.bias * self.lr_mul)\n        return out\n\n    def __repr__(self):\n        return f\"{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]})\"\n\n\nclass ScaledLeakyReLU(nn.Module):\n    def __init__(self, negative_slope=0.2):\n        super().__init__()\n\n        self.negative_slope = negative_slope\n\n    def forward(self, inputs):\n        out = F.leaky_relu(inputs, negative_slope=self.negative_slope)\n\n        return out * math.sqrt(2)\n\n\nclass ModulatedConv2d(nn.Module):\n    def __init__(\n        self,\n        in_channel,\n        out_channel,\n        kernel_size,\n        style_dim,\n        demodulate=True,\n        upsample=False,\n        downsample=False,\n        blur_kernel=[1, 3, 3, 1],\n    ):\n        super().__init__()\n\n        self.eps = 1e-8\n        self.kernel_size = kernel_size\n        self.in_channel = in_channel\n        self.out_channel = out_channel\n        self.upsample = upsample\n        self.downsample = downsample\n\n        if upsample:\n            factor = 2\n            p = (len(blur_kernel) - factor) - (kernel_size - 1)\n            pad0 = (p + 1) // 2 + factor - 1\n            pad1 = p // 2 + 1\n\n            self.blur = Blur(blur_kernel, pad=(pad0, pad1), upsample_factor=factor)\n\n        if downsample:\n            factor = 2\n            p = (len(blur_kernel) - factor) + (kernel_size - 1)\n            pad0 = (p + 1) // 2\n            pad1 = p // 2\n\n            self.blur = Blur(blur_kernel, pad=(pad0, pad1))\n\n        fan_in = in_channel * kernel_size ** 2\n        self.scale = 1 / math.sqrt(fan_in)\n        self.padding = kernel_size // 2\n\n        self.weight = nn.Parameter(th.randn(1, out_channel, in_channel, kernel_size, kernel_size))\n\n        self.modulation = EqualLinear(style_dim, in_channel, bias_init=1)\n\n        self.demodulate = demodulate\n\n    def __repr__(self):\n        return (\n            f\"{self.__class__.__name__}({self.in_channel}, {self.out_channel}, {self.kernel_size}, \"\n            f\"upsample={self.upsample}, downsample={self.downsample})\"\n        )\n\n    def forward(self, inputs, style):\n        batch, in_channel, height, width = inputs.shape\n\n        style = self.modulation(style).view(batch, 1, in_channel, 1, 1)\n        weight = self.scale * self.weight * style\n\n        if self.demodulate:\n            demod = th.rsqrt(weight.pow(2).sum([2, 3, 4]) + 1e-8)\n            weight = weight * demod.view(batch, self.out_channel, 1, 1, 1)\n\n        weight = weight.view(batch * self.out_channel, in_channel, self.kernel_size, self.kernel_size)\n\n        if self.upsample:\n            inputs = inputs.view(1, batch * in_channel, height, width)\n            weight = weight.view(batch, self.out_channel, in_channel, self.kernel_size, self.kernel_size)\n            weight = weight.transpose(1, 2).reshape(\n                batch * in_channel, self.out_channel, self.kernel_size, self.kernel_size\n            )\n            out = F.conv_transpose2d(inputs, weight, padding=0, stride=2, groups=batch)\n            _, _, height, width = out.shape\n            out = out.view(batch, self.out_channel, height, width)\n            out = self.blur(out)\n\n        elif self.downsample:\n            inputs = self.blur(inputs)\n            _, _, height, width = inputs.shape\n            inputs = inputs.view(1, batch * in_channel, height, width)\n            out = F.conv2d(inputs, weight, padding=0, stride=2, groups=batch)\n            _, _, height, width = out.shape\n            out = out.view(batch, self.out_channel, height, width)\n\n        else:\n            inputs = inputs.view(1, batch * in_channel, height, width)\n            out = F.conv2d(inputs, weight, padding=self.padding, groups=batch)\n            _, _, height, width = out.shape\n            out = out.view(batch, self.out_channel, height, width)\n\n        return out\n\n\nclass NoiseInjection(nn.Module):\n    def __init__(self):\n        super().__init__()\n        self.weight = nn.Parameter(th.zeros(1))\n\n    def forward(self, image, noise=None):\n        if noise is None:\n            batch, _, height, width = image.shape\n            noise = image.new_empty(batch, 1, height, width).normal_()\n        return image + self.weight * noise.to(image.device)\n\n\nclass ConstantInput(nn.Module):\n    def __init__(self, channel, size=4):\n        super().__init__()\n\n        self.input = nn.Parameter(th.randn(1, channel, size, size))\n\n    def forward(self, inputs):\n        batch = inputs.shape[0]\n        out = self.input.repeat(batch, 1, 1, 1)\n        return out\n\n\nclass LatentInput(nn.Module):\n    def __init__(self, latent_dim, channel, size=4):\n        super().__init__()\n        self.channel = channel\n        self.size = size\n        self.linear = EqualLinear(latent_dim, channel * size * size, activation=\"fused_lrelu\")\n        self.activate = FusedLeakyReLU(channel * size * size)\n        self.input = nn.Parameter(th.randn(1))\n\n    def forward(self, inputs):\n        batch = inputs.shape[0]\n        out = self.linear(inputs[:, 0])\n        out = self.activate(out)\n        return out.reshape((batch, self.channel, self.size, self.size))\n\n\nclass ManipulationLayer(th.nn.Module):\n    def __init__(self, layer):\n        super().__init__()\n        self.layer = layer\n\n    def forward(self, input, tranforms_dict_list):\n        out = input\n        for transform_dict in tranforms_dict_list:\n            if transform_dict[\"layer\"] == self.layer:\n                out = transform_dict[\"transform\"].to(out.device)(out)\n        return out\n\n\nclass StyledConv(nn.Module):\n    def __init__(\n        self,\n        in_channel,\n        out_channel,\n        kernel_size,\n        style_dim,\n        upsample=False,\n        blur_kernel=[1, 3, 3, 1],\n        demodulate=True,\n        layerID=-1,\n    ):\n        super().__init__()\n\n        self.conv = ModulatedConv2d(\n            in_channel,\n            out_channel,\n            kernel_size,\n            style_dim,\n            upsample=upsample,\n            blur_kernel=blur_kernel,\n            demodulate=demodulate,\n        )\n\n        self.noise = NoiseInjection()\n        self.activate = FusedLeakyReLU(out_channel)\n        self.manipulation = ManipulationLayer(layerID)\n\n    def forward(self, inputs, style, noise=None, transform_dict_list=[]):\n        out = self.conv(inputs, style)\n        out = self.noise(out, noise=noise)\n        out = self.activate(out)\n        out = self.manipulation(out, transform_dict_list)\n        return out\n\n\nclass ToRGB(nn.Module):\n    def __init__(self, in_channel, style_dim, upsample=True, blur_kernel=[1, 3, 3, 1]):\n        super().__init__()\n\n        if upsample:\n            self.upsample = Upsample(blur_kernel)\n\n        self.conv = ModulatedConv2d(in_channel, 3, 1, style_dim, demodulate=False)\n        self.bias = nn.Parameter(th.zeros(1, 3, 1, 1))\n\n    def forward(self, inputs, style, skip=None):\n        out = self.conv(inputs, style)\n        out = out + self.bias\n\n        if skip is not None:\n            skip = self.upsample(skip)\n\n            out = out + skip\n\n        return out\n\n\nclass Generator(nn.Module):\n    def __init__(\n        self,\n        size,\n        style_dim,\n        n_mlp,\n        channel_multiplier=2,\n        blur_kernel=[1, 3, 3, 1],\n        lr_mlp=0.01,\n        constant_input=False,\n        checkpoint=None,\n        output_size=None,\n        min_rgb_size=4,\n        base_res_factor=1,\n    ):\n        super().__init__()\n\n        self.size = size\n        self.style_dim = style_dim\n\n        layers = [PixelNorm()]\n\n        for i in range(n_mlp):\n            layers.append(EqualLinear(style_dim, style_dim, lr_mul=lr_mlp, activation=\"fused_lrelu\"))\n\n        self.style = nn.Sequential(*layers)\n\n        self.channels = {\n            4: 512,\n            8: 512,\n            16: 512,\n            32: 512,\n            64: 256 * channel_multiplier,\n            128: 128 * channel_multiplier,\n            256: 64 * channel_multiplier,\n            512: 32 * channel_multiplier,\n            1024: 16 * channel_multiplier,\n        }\n\n        self.log_size = int(math.log(size, 2))\n        self.num_layers = (self.log_size - 2) * 2 + 1\n        self.n_latent = self.log_size * 2 - 2\n        self.min_rgb_size = min_rgb_size\n\n        if constant_input:\n            self.input = ConstantInput(self.channels[4])\n        else:\n            self.input = LatentInput(style_dim, self.channels[4])\n\n        self.const_manipulation = ManipulationLayer(0)\n\n        layerID = 1\n        self.conv1 = StyledConv(\n            self.channels[4], self.channels[4], 3, style_dim, blur_kernel=blur_kernel, layerID=layerID\n        )\n        self.to_rgb1 = ToRGB(self.channels[4], style_dim, upsample=False)\n\n        self.convs = nn.ModuleList()\n        self.upsamples = nn.ModuleList()\n        self.to_rgbs = nn.ModuleList()\n        self.noises = nn.Module()\n\n        in_channel = self.channels[4]\n\n        for layer_idx in range(self.num_layers):\n            res = (layer_idx + 5) // 2\n            shape = [1, 1, 2 ** res, 2 ** res]\n            self.noises.register_buffer(f\"noise_{layer_idx}\", th.randn(*shape))\n\n        for i in range(3, self.log_size + 1):\n            out_channel = self.channels[2 ** i]\n\n            layerID += 1\n            self.convs.append(\n                StyledConv(\n                    in_channel, out_channel, 3, style_dim, upsample=True, blur_kernel=blur_kernel, layerID=layerID\n                )\n            )\n\n            layerID += 1\n            self.convs.append(\n                StyledConv(out_channel, out_channel, 3, style_dim, blur_kernel=blur_kernel, layerID=layerID)\n            )\n\n            self.to_rgbs.append(ToRGB(out_channel, style_dim))\n\n            in_channel = out_channel\n\n        self.truncation_latent = None\n\n        if checkpoint is not None:\n            self.load_state_dict(th.load(checkpoint)[\"g_ema\"])\n\n        if size != output_size or base_res_factor != 1:\n            for layer_idx in range(self.num_layers):\n                res = (layer_idx + 5) // 2\n                shape = [\n                    1,\n                    1,\n                    int(base_res_factor * 2 ** res * (2 if output_size == 1080 else 1)),\n                    int(base_res_factor * 2 ** res * (2 if output_size == 1920 else 1)),\n                ]\n                setattr(self.noises, f\"noise_{layer_idx}\", th.randn(*shape))\n\n    def make_noise(self):\n        device = self.input.input.device\n\n        noises = [th.randn(1, 1, 2 ** 2, 2 ** 2, device=device)]\n\n        for i in range(3, self.log_size + 1):\n            for _ in range(2):\n                noises.append(th.randn(1, 1, 2 ** i, 2 ** i, device=device))\n\n        return noises\n\n    def mean_latent(self, n_latent):\n        latent_in = th.randn(n_latent, self.style_dim, device=self.input.input.device)\n        latent = self.style(latent_in).mean(0, keepdim=True)\n\n        return latent\n\n    def get_latent(self, inputs):\n        return self.style(inputs)\n\n    def forward(\n        self,\n        styles,\n        return_latents=False,\n        return_activation_maps=False,\n        inject_index=None,\n        truncation=1.0,\n        truncation_latent=None,\n        input_is_latent=False,\n        noise=None,\n        randomize_noise=True,\n        transform_dict_list=[],\n        map_latents=False,\n    ):\n        if map_latents:\n            latent = th.cat([self.style(s[None, None, :]) for s in styles], axis=0)\n            latent = latent.repeat(1, self.n_latent, 1)\n            return latent\n\n        if not input_is_latent:\n            styles = [self.style(s) for s in styles]\n\n            if len(styles) < 2:\n                inject_index = self.n_latent\n                if styles[0].ndim < 3:\n                    latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1)\n                else:\n                    latent = styles[0]\n            else:\n                if inject_index is None:\n                    inject_index = random.randint(1, self.n_latent - 1)\n                latent = styles[0].unsqueeze(1).repeat(1, inject_index, 1)\n                latent2 = styles[1].unsqueeze(1).repeat(1, self.n_latent - inject_index, 1)\n                latent = th.cat([latent, latent2], 1)\n        else:\n            latent = styles\n            if latent.dim() == 2:\n                latent = latent[:, None, :].repeat(1, self.n_latent, 1)\n\n        if noise is None:\n            noise = [None] * self.num_layers\n        for ns, noise_scale in enumerate(noise):\n            if not randomize_noise and noise_scale is None:\n                noise[ns] = getattr(self.noises, f\"noise_{ns}\")\n\n        if isinstance(truncation, float):\n            truncation = th.cuda.FloatTensor([truncation])\n        if self.truncation_latent is None:\n            self.truncation_latent = truncation_latent if truncation_latent is not None else self.mean_latent(2 ** 14)\n        latent = self.truncation_latent[None, ...] + truncation.to(latent.device)[:, None, None] * (\n            latent - self.truncation_latent[None, ...]\n        )\n\n        activation_map_list = []\n\n        out = self.input(latent)\n        out = self.const_manipulation(out, transform_dict_list)\n        out = self.conv1(out, latent[:, 0], noise=noise[0], transform_dict_list=transform_dict_list)\n        activation_map_list.append(out)\n\n        current_size = 4\n        if self.min_rgb_size <= current_size:\n            image = self.to_rgb1(out, latent[:, 1])\n        else:\n            image = None\n\n        i = 1\n        for conv1, conv2, noise1, noise2, to_rgb in zip(\n            self.convs[::2], self.convs[1::2], noise[1::2], noise[2::2], self.to_rgbs\n        ):\n            out = conv1(out, latent[:, i], noise=noise1, transform_dict_list=transform_dict_list)\n            current_size *= 2\n            activation_map_list.append(out)\n            out = conv2(out, latent[:, i + 1], noise=noise2, transform_dict_list=transform_dict_list)\n            activation_map_list.append(out)\n            if self.min_rgb_size <= current_size:\n                image = to_rgb(out, latent[:, i + 2], image)\n            i += 2\n\n        if return_activation_maps:\n            return image, activation_map_list\n        elif return_latents:\n            return image, latent\n        else:\n            return image, None\n\n\nclass ConvLayer(nn.Sequential):\n    def __init__(\n        self,\n        in_channel,\n        out_channel,\n        kernel_size,\n        downsample=False,\n        blur_kernel=[1, 3, 3, 1],\n        bias=True,\n        activate=True,\n    ):\n        layers = []\n\n        if downsample:\n            factor = 2\n            p = (len(blur_kernel) - factor) + (kernel_size - 1)\n            pad0 = (p + 1) // 2\n            pad1 = p // 2\n\n            layers.append(Blur(blur_kernel, pad=(pad0, pad1)))\n\n            stride = 2\n            self.padding = 0\n\n        else:\n            stride = 1\n            self.padding = kernel_size // 2\n\n        layers.append(\n            EqualConv2d(\n                in_channel, out_channel, kernel_size, padding=self.padding, stride=stride, bias=bias and not activate,\n            )\n        )\n\n        if activate:\n            if bias:\n                layers.append(FusedLeakyReLU(out_channel))\n\n            else:\n                layers.append(ScaledLeakyReLU(0.2))\n\n        super().__init__(*layers)\n\n\nclass ResBlock(nn.Module):\n    def __init__(self, in_channel, out_channel, blur_kernel=[1, 3, 3, 1], use_skip=True):\n        super().__init__()\n\n        self.conv1 = ConvLayer(in_channel, in_channel, 3)\n        self.conv2 = ConvLayer(in_channel, out_channel, 3, downsample=True)\n\n        if use_skip:\n            self.skip = ConvLayer(in_channel, out_channel, 1, downsample=True, activate=False, bias=False)\n        else:\n            self.skip = None\n\n    def forward(self, inputs):\n        out = self.conv1(inputs)\n        out = self.conv2(out)\n\n        if self.skip is not None:\n            skip = self.skip(inputs)\n            out = (out + skip) / math.sqrt(2)\n\n        return out\n\n\nclass Discriminator(nn.Module):\n    def __init__(self, size, channel_multiplier=2, blur_kernel=[1, 3, 3, 1], use_skip=True):\n        super().__init__()\n\n        channels = {\n            4: 512,\n            8: 512,\n            16: 512,\n            32: 512,\n            64: 256 * channel_multiplier,\n            128: 128 * channel_multiplier,\n            256: 64 * channel_multiplier,\n            512: 32 * channel_multiplier,\n            1024: 16 * channel_multiplier,\n        }\n\n        convs = [ConvLayer(3, channels[size], 1)]\n\n        log_size = int(math.log(size, 2))\n\n        in_channel = channels[size]\n\n        for i in range(log_size, 2, -1):\n            out_channel = channels[2 ** (i - 1)]\n\n            convs.append(ResBlock(in_channel, out_channel, blur_kernel, use_skip=use_skip))\n\n            in_channel = out_channel\n\n        self.convs = nn.Sequential(*convs)\n\n        self.stddev_group = 4\n        self.stddev_feat = 1\n\n        self.final_conv = ConvLayer(in_channel + 1, channels[4], 3)\n        self.final_linear = nn.Sequential(\n            EqualLinear(channels[4] * 4 * 4, channels[4], activation=\"fused_lrelu\"), EqualLinear(channels[4], 1),\n        )\n\n    def forward(self, inputs):\n        out = self.convs(inputs)\n\n        batch, channel, height, width = out.shape\n\n        try:\n            group = min(batch, self.stddev_group)\n            stddev = out.view(group, -1, self.stddev_feat, channel // self.stddev_feat, height, width)\n            stddev = th.sqrt(stddev.var(0, unbiased=False) + 1e-8)\n            stddev = stddev.mean([2, 3, 4], keepdims=True).squeeze(2)\n            stddev = stddev.repeat(group, 1, height, width)\n            out = th.cat([out, stddev], 1)\n        except RuntimeError:\n            group = batch\n            stddev = out.view(group, -1, self.stddev_feat, channel // self.stddev_feat, height, width)\n            stddev = th.sqrt(stddev.var(0, unbiased=False) + 1e-8)\n            stddev = stddev.mean([2, 3, 4], keepdims=True).squeeze(2)\n            stddev = stddev.repeat(group, 1, height, width)\n            out = th.cat([out, stddev], 1)\n\n        out = self.final_conv(out)\n\n        out = out.view(batch, -1)\n        out = self.final_linear(out)\n\n        return out\n"
  },
  {
    "path": "op/__init__.py",
    "content": "from .fused_act import FusedLeakyReLU, fused_leaky_relu\nfrom .upfirdn2d import upfirdn2d\n"
  },
  {
    "path": "op/fused_act.py",
    "content": "import os\r\n\r\nimport torch\r\nfrom torch import nn\r\nfrom torch.nn import functional as F\r\nfrom torch.autograd import Function\r\nfrom torch.utils.cpp_extension import load\r\n\r\n\r\nmodule_path = os.path.dirname(__file__)\r\nfused = load(\r\n    \"fused\",\r\n    sources=[\r\n        os.path.join(module_path, \"fused_bias_act.cpp\"),\r\n        os.path.join(module_path, \"fused_bias_act_kernel.cu\"),\r\n    ],\r\n)\r\n\r\n\r\nclass FusedLeakyReLUFunctionBackward(Function):\r\n    @staticmethod\r\n    def forward(ctx, grad_output, out, negative_slope, scale):\r\n        ctx.save_for_backward(out)\r\n        ctx.negative_slope = negative_slope\r\n        ctx.scale = scale\r\n\r\n        empty = grad_output.new_empty(0)\r\n\r\n        grad_input = fused.fused_bias_act(\r\n            grad_output, empty, out, 3, 1, negative_slope, scale\r\n        )\r\n\r\n        dim = [0]\r\n\r\n        if grad_input.ndim > 2:\r\n            dim += list(range(2, grad_input.ndim))\r\n\r\n        grad_bias = grad_input.sum(dim).detach()\r\n\r\n        return grad_input, grad_bias\r\n\r\n    @staticmethod\r\n    def backward(ctx, gradgrad_input, gradgrad_bias):\r\n        out, = ctx.saved_tensors\r\n        gradgrad_out = fused.fused_bias_act(\r\n            gradgrad_input, gradgrad_bias, out, 3, 1, ctx.negative_slope, ctx.scale\r\n        )\r\n\r\n        return gradgrad_out, None, None, None\r\n\r\n\r\nclass FusedLeakyReLUFunction(Function):\r\n    @staticmethod\r\n    def forward(ctx, input, bias, negative_slope, scale):\r\n        empty = input.new_empty(0)\r\n        out = fused.fused_bias_act(input, bias, empty, 3, 0, negative_slope, scale)\r\n        ctx.save_for_backward(out)\r\n        ctx.negative_slope = negative_slope\r\n        ctx.scale = scale\r\n\r\n        return out\r\n\r\n    @staticmethod\r\n    def backward(ctx, grad_output):\r\n        out, = ctx.saved_tensors\r\n\r\n        grad_input, grad_bias = FusedLeakyReLUFunctionBackward.apply(\r\n            grad_output, out, ctx.negative_slope, ctx.scale\r\n        )\r\n\r\n        return grad_input, grad_bias, None, None\r\n\r\n\r\nclass FusedLeakyReLU(nn.Module):\r\n    def __init__(self, channel, negative_slope=0.2, scale=2 ** 0.5):\r\n        super().__init__()\r\n\r\n        self.bias = nn.Parameter(torch.zeros(channel))\r\n        self.negative_slope = negative_slope\r\n        self.scale = scale\r\n\r\n    def forward(self, input):\r\n        return fused_leaky_relu(input, self.bias, self.negative_slope, self.scale)\r\n\r\n\r\ndef fused_leaky_relu(input, bias, negative_slope=0.2, scale=2 ** 0.5):\r\n    if input.device.type == \"cpu\":\r\n        rest_dim = [1] * (input.ndim - bias.ndim - 1)\r\n        return (\r\n            F.leaky_relu(\r\n                input + bias.view(1, bias.shape[0], *rest_dim), negative_slope=0.2\r\n            )\r\n            * scale\r\n        )\r\n\r\n    else:\r\n        return FusedLeakyReLUFunction.apply(input, bias, negative_slope, scale)\r\n"
  },
  {
    "path": "op/fused_bias_act.cpp",
    "content": "#include <torch/extension.h>\r\n\r\n\r\ntorch::Tensor fused_bias_act_op(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer,\r\n    int act, int grad, float alpha, float scale);\r\n\r\n#define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x \" must be a CUDA tensor\")\r\n#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x \" must be contiguous\")\r\n#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)\r\n\r\ntorch::Tensor fused_bias_act(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer,\r\n    int act, int grad, float alpha, float scale) {\r\n    CHECK_CUDA(input);\r\n    CHECK_CUDA(bias);\r\n\r\n    return fused_bias_act_op(input, bias, refer, act, grad, alpha, scale);\r\n}\r\n\r\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\r\n    m.def(\"fused_bias_act\", &fused_bias_act, \"fused bias act (CUDA)\");\r\n}"
  },
  {
    "path": "op/fused_bias_act_kernel.cu",
    "content": "// Copyright (c) 2019, NVIDIA Corporation. All rights reserved.\r\n//\r\n// This work is made available under the Nvidia Source Code License-NC.\r\n// To view a copy of this license, visit\r\n// https://nvlabs.github.io/stylegan2/license.html\r\n\r\n#include <torch/types.h>\r\n\r\n#include <ATen/ATen.h>\r\n#include <ATen/AccumulateType.h>\r\n#include <ATen/cuda/CUDAContext.h>\r\n#include <ATen/cuda/CUDAApplyUtils.cuh>\r\n\r\n#include <cuda.h>\r\n#include <cuda_runtime.h>\r\n\r\n\r\ntemplate <typename scalar_t>\r\nstatic __global__ void fused_bias_act_kernel(scalar_t* out, const scalar_t* p_x, const scalar_t* p_b, const scalar_t* p_ref,\r\n    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) {\r\n    int xi = blockIdx.x * loop_x * blockDim.x + threadIdx.x;\r\n\r\n    scalar_t zero = 0.0;\r\n\r\n    for (int loop_idx = 0; loop_idx < loop_x && xi < size_x; loop_idx++, xi += blockDim.x) {\r\n        scalar_t x = p_x[xi];\r\n\r\n        if (use_bias) {\r\n            x += p_b[(xi / step_b) % size_b];\r\n        }\r\n\r\n        scalar_t ref = use_ref ? p_ref[xi] : zero;\r\n\r\n        scalar_t y;\r\n\r\n        switch (act * 10 + grad) {\r\n            default:\r\n            case 10: y = x; break;\r\n            case 11: y = x; break;\r\n            case 12: y = 0.0; break;\r\n\r\n            case 30: y = (x > 0.0) ? x : x * alpha; break;\r\n            case 31: y = (ref > 0.0) ? x : x * alpha; break;\r\n            case 32: y = 0.0; break;\r\n        }\r\n\r\n        out[xi] = y * scale;\r\n    }\r\n}\r\n\r\n\r\ntorch::Tensor fused_bias_act_op(const torch::Tensor& input, const torch::Tensor& bias, const torch::Tensor& refer,\r\n    int act, int grad, float alpha, float scale) {\r\n    int curDevice = -1;\r\n    cudaGetDevice(&curDevice);\r\n    cudaStream_t stream = at::cuda::getCurrentCUDAStream(curDevice);\r\n\r\n    auto x = input.contiguous();\r\n    auto b = bias.contiguous();\r\n    auto ref = refer.contiguous();\r\n\r\n    int use_bias = b.numel() ? 1 : 0;\r\n    int use_ref = ref.numel() ? 1 : 0;\r\n\r\n    int size_x = x.numel();\r\n    int size_b = b.numel();\r\n    int step_b = 1;\r\n\r\n    for (int i = 1 + 1; i < x.dim(); i++) {\r\n        step_b *= x.size(i);\r\n    }\r\n\r\n    int loop_x = 4;\r\n    int block_size = 4 * 32;\r\n    int grid_size = (size_x - 1) / (loop_x * block_size) + 1;\r\n\r\n    auto y = torch::empty_like(x);\r\n\r\n    AT_DISPATCH_FLOATING_TYPES_AND_HALF(x.scalar_type(), \"fused_bias_act_kernel\", [&] {\r\n        fused_bias_act_kernel<scalar_t><<<grid_size, block_size, 0, stream>>>(\r\n            y.data_ptr<scalar_t>(),\r\n            x.data_ptr<scalar_t>(),\r\n            b.data_ptr<scalar_t>(),\r\n            ref.data_ptr<scalar_t>(),\r\n            act,\r\n            grad,\r\n            alpha,\r\n            scale,\r\n            loop_x,\r\n            size_x,\r\n            step_b,\r\n            size_b,\r\n            use_bias,\r\n            use_ref\r\n        );\r\n    });\r\n\r\n    return y;\r\n}"
  },
  {
    "path": "op/upfirdn2d.cpp",
    "content": "#include <torch/extension.h>\r\n\r\n\r\ntorch::Tensor upfirdn2d_op(const torch::Tensor& input, const torch::Tensor& kernel,\r\n                            int up_x, int up_y, int down_x, int down_y,\r\n                            int pad_x0, int pad_x1, int pad_y0, int pad_y1);\r\n\r\n#define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x \" must be a CUDA tensor\")\r\n#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x \" must be contiguous\")\r\n#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)\r\n\r\ntorch::Tensor upfirdn2d(const torch::Tensor& input, const torch::Tensor& kernel,\r\n                        int up_x, int up_y, int down_x, int down_y,\r\n                        int pad_x0, int pad_x1, int pad_y0, int pad_y1) {\r\n    CHECK_CUDA(input);\r\n    CHECK_CUDA(kernel);\r\n\r\n    return upfirdn2d_op(input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1);\r\n}\r\n\r\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\r\n    m.def(\"upfirdn2d\", &upfirdn2d, \"upfirdn2d (CUDA)\");\r\n}"
  },
  {
    "path": "op/upfirdn2d.py",
    "content": "import os\r\n\r\nimport torch\r\nfrom torch.nn import functional as F\r\nfrom torch.autograd import Function\r\nfrom torch.utils.cpp_extension import load\r\n\r\n\r\nmodule_path = os.path.dirname(__file__)\r\nupfirdn2d_op = load(\r\n    \"upfirdn2d\",\r\n    sources=[\r\n        os.path.join(module_path, \"upfirdn2d.cpp\"),\r\n        os.path.join(module_path, \"upfirdn2d_kernel.cu\"),\r\n    ],\r\n)\r\n\r\n\r\nclass UpFirDn2dBackward(Function):\r\n    @staticmethod\r\n    def forward(\r\n        ctx, grad_output, kernel, grad_kernel, up, down, pad, g_pad, in_size, out_size\r\n    ):\r\n\r\n        up_x, up_y = up\r\n        down_x, down_y = down\r\n        g_pad_x0, g_pad_x1, g_pad_y0, g_pad_y1 = g_pad\r\n\r\n        grad_output = grad_output.reshape(-1, out_size[0], out_size[1], 1)\r\n\r\n        grad_input = upfirdn2d_op.upfirdn2d(\r\n            grad_output,\r\n            grad_kernel,\r\n            down_x,\r\n            down_y,\r\n            up_x,\r\n            up_y,\r\n            g_pad_x0,\r\n            g_pad_x1,\r\n            g_pad_y0,\r\n            g_pad_y1,\r\n        )\r\n        grad_input = grad_input.view(in_size[0], in_size[1], in_size[2], in_size[3])\r\n\r\n        ctx.save_for_backward(kernel)\r\n\r\n        pad_x0, pad_x1, pad_y0, pad_y1 = pad\r\n\r\n        ctx.up_x = up_x\r\n        ctx.up_y = up_y\r\n        ctx.down_x = down_x\r\n        ctx.down_y = down_y\r\n        ctx.pad_x0 = pad_x0\r\n        ctx.pad_x1 = pad_x1\r\n        ctx.pad_y0 = pad_y0\r\n        ctx.pad_y1 = pad_y1\r\n        ctx.in_size = in_size\r\n        ctx.out_size = out_size\r\n\r\n        return grad_input\r\n\r\n    @staticmethod\r\n    def backward(ctx, gradgrad_input):\r\n        kernel, = ctx.saved_tensors\r\n\r\n        gradgrad_input = gradgrad_input.reshape(-1, ctx.in_size[2], ctx.in_size[3], 1)\r\n\r\n        gradgrad_out = upfirdn2d_op.upfirdn2d(\r\n            gradgrad_input,\r\n            kernel,\r\n            ctx.up_x,\r\n            ctx.up_y,\r\n            ctx.down_x,\r\n            ctx.down_y,\r\n            ctx.pad_x0,\r\n            ctx.pad_x1,\r\n            ctx.pad_y0,\r\n            ctx.pad_y1,\r\n        )\r\n        # gradgrad_out = gradgrad_out.view(ctx.in_size[0], ctx.out_size[0], ctx.out_size[1], ctx.in_size[3])\r\n        gradgrad_out = gradgrad_out.view(\r\n            ctx.in_size[0], ctx.in_size[1], ctx.out_size[0], ctx.out_size[1]\r\n        )\r\n\r\n        return gradgrad_out, None, None, None, None, None, None, None, None\r\n\r\n\r\nclass UpFirDn2d(Function):\r\n    @staticmethod\r\n    def forward(ctx, input, kernel, up, down, pad):\r\n        up_x, up_y = up\r\n        down_x, down_y = down\r\n        pad_x0, pad_x1, pad_y0, pad_y1 = pad\r\n\r\n        kernel_h, kernel_w = kernel.shape\r\n        batch, channel, in_h, in_w = input.shape\r\n        ctx.in_size = input.shape\r\n\r\n        input = input.reshape(-1, in_h, in_w, 1)\r\n\r\n        ctx.save_for_backward(kernel, torch.flip(kernel, [0, 1]))\r\n\r\n        out_h = (in_h * up_y + pad_y0 + pad_y1 - kernel_h) // down_y + 1\r\n        out_w = (in_w * up_x + pad_x0 + pad_x1 - kernel_w) // down_x + 1\r\n        ctx.out_size = (out_h, out_w)\r\n\r\n        ctx.up = (up_x, up_y)\r\n        ctx.down = (down_x, down_y)\r\n        ctx.pad = (pad_x0, pad_x1, pad_y0, pad_y1)\r\n\r\n        g_pad_x0 = kernel_w - pad_x0 - 1\r\n        g_pad_y0 = kernel_h - pad_y0 - 1\r\n        g_pad_x1 = in_w * up_x - out_w * down_x + pad_x0 - up_x + 1\r\n        g_pad_y1 = in_h * up_y - out_h * down_y + pad_y0 - up_y + 1\r\n\r\n        ctx.g_pad = (g_pad_x0, g_pad_x1, g_pad_y0, g_pad_y1)\r\n\r\n        out = upfirdn2d_op.upfirdn2d(\r\n            input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1\r\n        )\r\n        # out = out.view(major, out_h, out_w, minor)\r\n        out = out.view(-1, channel, out_h, out_w)\r\n\r\n        return out\r\n\r\n    @staticmethod\r\n    def backward(ctx, grad_output):\r\n        kernel, grad_kernel = ctx.saved_tensors\r\n\r\n        grad_input = UpFirDn2dBackward.apply(\r\n            grad_output,\r\n            kernel,\r\n            grad_kernel,\r\n            ctx.up,\r\n            ctx.down,\r\n            ctx.pad,\r\n            ctx.g_pad,\r\n            ctx.in_size,\r\n            ctx.out_size,\r\n        )\r\n\r\n        return grad_input, None, None, None, None\r\n\r\n\r\ndef upfirdn2d(input, kernel, up=1, down=1, pad=(0, 0)):\r\n    if input.device.type == \"cpu\":\r\n        out = upfirdn2d_native(\r\n            input, kernel, up, up, down, down, pad[0], pad[1], pad[0], pad[1]\r\n        )\r\n\r\n    else:\r\n        out = UpFirDn2d.apply(\r\n            input, kernel, (up, up), (down, down), (pad[0], pad[1], pad[0], pad[1])\r\n        )\r\n\r\n    return out\r\n\r\n\r\ndef upfirdn2d_native(\r\n    input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1\r\n):\r\n    _, channel, in_h, in_w = input.shape\r\n    input = input.reshape(-1, in_h, in_w, 1)\r\n\r\n    _, in_h, in_w, minor = input.shape\r\n    kernel_h, kernel_w = kernel.shape\r\n\r\n    out = input.view(-1, in_h, 1, in_w, 1, minor)\r\n    out = F.pad(out, [0, 0, 0, up_x - 1, 0, 0, 0, up_y - 1])\r\n    out = out.view(-1, in_h * up_y, in_w * up_x, minor)\r\n\r\n    out = F.pad(\r\n        out, [0, 0, max(pad_x0, 0), max(pad_x1, 0), max(pad_y0, 0), max(pad_y1, 0)]\r\n    )\r\n    out = out[\r\n        :,\r\n        max(-pad_y0, 0) : out.shape[1] - max(-pad_y1, 0),\r\n        max(-pad_x0, 0) : out.shape[2] - max(-pad_x1, 0),\r\n        :,\r\n    ]\r\n\r\n    out = out.permute(0, 3, 1, 2)\r\n    out = out.reshape(\r\n        [-1, 1, in_h * up_y + pad_y0 + pad_y1, in_w * up_x + pad_x0 + pad_x1]\r\n    )\r\n    w = torch.flip(kernel, [0, 1]).view(1, 1, kernel_h, kernel_w)\r\n    out = F.conv2d(out, w)\r\n    out = out.reshape(\r\n        -1,\r\n        minor,\r\n        in_h * up_y + pad_y0 + pad_y1 - kernel_h + 1,\r\n        in_w * up_x + pad_x0 + pad_x1 - kernel_w + 1,\r\n    )\r\n    out = out.permute(0, 2, 3, 1)\r\n    out = out[:, ::down_y, ::down_x, :]\r\n\r\n    out_h = (in_h * up_y + pad_y0 + pad_y1 - kernel_h) // down_y + 1\r\n    out_w = (in_w * up_x + pad_x0 + pad_x1 - kernel_w) // down_x + 1\r\n\r\n    return out.view(-1, channel, out_h, out_w)\r\n"
  },
  {
    "path": "op/upfirdn2d_kernel.cu",
    "content": "// Copyright (c) 2019, NVIDIA Corporation. All rights reserved.\r\n//\r\n// This work is made available under the Nvidia Source Code License-NC.\r\n// To view a copy of this license, visit\r\n// https://nvlabs.github.io/stylegan2/license.html\r\n\r\n#include <torch/types.h>\r\n\r\n#include <ATen/ATen.h>\r\n#include <ATen/AccumulateType.h>\r\n#include <ATen/cuda/CUDAApplyUtils.cuh>\r\n#include <ATen/cuda/CUDAContext.h>\r\n\r\n#include <cuda.h>\r\n#include <cuda_runtime.h>\r\n\r\nstatic __host__ __device__ __forceinline__ int floor_div(int a, int b) {\r\n  int c = a / b;\r\n\r\n  if (c * b > a) {\r\n    c--;\r\n  }\r\n\r\n  return c;\r\n}\r\n\r\nstruct UpFirDn2DKernelParams {\r\n  int up_x;\r\n  int up_y;\r\n  int down_x;\r\n  int down_y;\r\n  int pad_x0;\r\n  int pad_x1;\r\n  int pad_y0;\r\n  int pad_y1;\r\n\r\n  int major_dim;\r\n  int in_h;\r\n  int in_w;\r\n  int minor_dim;\r\n  int kernel_h;\r\n  int kernel_w;\r\n  int out_h;\r\n  int out_w;\r\n  int loop_major;\r\n  int loop_x;\r\n};\r\n\r\ntemplate <typename scalar_t>\r\n__global__ void upfirdn2d_kernel_large(scalar_t *out, const scalar_t *input,\r\n                                       const scalar_t *kernel,\r\n                                       const UpFirDn2DKernelParams p) {\r\n  int minor_idx = blockIdx.x * blockDim.x + threadIdx.x;\r\n  int out_y = minor_idx / p.minor_dim;\r\n  minor_idx -= out_y * p.minor_dim;\r\n  int out_x_base = blockIdx.y * p.loop_x * blockDim.y + threadIdx.y;\r\n  int major_idx_base = blockIdx.z * p.loop_major;\r\n\r\n  if (out_x_base >= p.out_w || out_y >= p.out_h ||\r\n      major_idx_base >= p.major_dim) {\r\n    return;\r\n  }\r\n\r\n  int mid_y = out_y * p.down_y + p.up_y - 1 - p.pad_y0;\r\n  int in_y = min(max(floor_div(mid_y, p.up_y), 0), p.in_h);\r\n  int h = min(max(floor_div(mid_y + p.kernel_h, p.up_y), 0), p.in_h) - in_y;\r\n  int kernel_y = mid_y + p.kernel_h - (in_y + 1) * p.up_y;\r\n\r\n  for (int loop_major = 0, major_idx = major_idx_base;\r\n       loop_major < p.loop_major && major_idx < p.major_dim;\r\n       loop_major++, major_idx++) {\r\n    for (int loop_x = 0, out_x = out_x_base;\r\n         loop_x < p.loop_x && out_x < p.out_w; loop_x++, out_x += blockDim.y) {\r\n      int mid_x = out_x * p.down_x + p.up_x - 1 - p.pad_x0;\r\n      int in_x = min(max(floor_div(mid_x, p.up_x), 0), p.in_w);\r\n      int w = min(max(floor_div(mid_x + p.kernel_w, p.up_x), 0), p.in_w) - in_x;\r\n      int kernel_x = mid_x + p.kernel_w - (in_x + 1) * p.up_x;\r\n\r\n      const scalar_t *x_p =\r\n          &input[((major_idx * p.in_h + in_y) * p.in_w + in_x) * p.minor_dim +\r\n                 minor_idx];\r\n      const scalar_t *k_p = &kernel[kernel_y * p.kernel_w + kernel_x];\r\n      int x_px = p.minor_dim;\r\n      int k_px = -p.up_x;\r\n      int x_py = p.in_w * p.minor_dim;\r\n      int k_py = -p.up_y * p.kernel_w;\r\n\r\n      scalar_t v = 0.0f;\r\n\r\n      for (int y = 0; y < h; y++) {\r\n        for (int x = 0; x < w; x++) {\r\n          v += static_cast<scalar_t>(*x_p) * static_cast<scalar_t>(*k_p);\r\n          x_p += x_px;\r\n          k_p += k_px;\r\n        }\r\n\r\n        x_p += x_py - w * x_px;\r\n        k_p += k_py - w * k_px;\r\n      }\r\n\r\n      out[((major_idx * p.out_h + out_y) * p.out_w + out_x) * p.minor_dim +\r\n          minor_idx] = v;\r\n    }\r\n  }\r\n}\r\n\r\ntemplate <typename scalar_t, int up_x, int up_y, int down_x, int down_y,\r\n          int kernel_h, int kernel_w, int tile_out_h, int tile_out_w>\r\n__global__ void upfirdn2d_kernel(scalar_t *out, const scalar_t *input,\r\n                                 const scalar_t *kernel,\r\n                                 const UpFirDn2DKernelParams p) {\r\n  const int tile_in_h = ((tile_out_h - 1) * down_y + kernel_h - 1) / up_y + 1;\r\n  const int tile_in_w = ((tile_out_w - 1) * down_x + kernel_w - 1) / up_x + 1;\r\n\r\n  __shared__ volatile float sk[kernel_h][kernel_w];\r\n  __shared__ volatile float sx[tile_in_h][tile_in_w];\r\n\r\n  int minor_idx = blockIdx.x;\r\n  int tile_out_y = minor_idx / p.minor_dim;\r\n  minor_idx -= tile_out_y * p.minor_dim;\r\n  tile_out_y *= tile_out_h;\r\n  int tile_out_x_base = blockIdx.y * p.loop_x * tile_out_w;\r\n  int major_idx_base = blockIdx.z * p.loop_major;\r\n\r\n  if (tile_out_x_base >= p.out_w | tile_out_y >= p.out_h |\r\n      major_idx_base >= p.major_dim) {\r\n    return;\r\n  }\r\n\r\n  for (int tap_idx = threadIdx.x; tap_idx < kernel_h * kernel_w;\r\n       tap_idx += blockDim.x) {\r\n    int ky = tap_idx / kernel_w;\r\n    int kx = tap_idx - ky * kernel_w;\r\n    scalar_t v = 0.0;\r\n\r\n    if (kx < p.kernel_w & ky < p.kernel_h) {\r\n      v = kernel[(p.kernel_h - 1 - ky) * p.kernel_w + (p.kernel_w - 1 - kx)];\r\n    }\r\n\r\n    sk[ky][kx] = v;\r\n  }\r\n\r\n  for (int loop_major = 0, major_idx = major_idx_base;\r\n       loop_major < p.loop_major & major_idx < p.major_dim;\r\n       loop_major++, major_idx++) {\r\n    for (int loop_x = 0, tile_out_x = tile_out_x_base;\r\n         loop_x < p.loop_x & tile_out_x < p.out_w;\r\n         loop_x++, tile_out_x += tile_out_w) {\r\n      int tile_mid_x = tile_out_x * down_x + up_x - 1 - p.pad_x0;\r\n      int tile_mid_y = tile_out_y * down_y + up_y - 1 - p.pad_y0;\r\n      int tile_in_x = floor_div(tile_mid_x, up_x);\r\n      int tile_in_y = floor_div(tile_mid_y, up_y);\r\n\r\n      __syncthreads();\r\n\r\n      for (int in_idx = threadIdx.x; in_idx < tile_in_h * tile_in_w;\r\n           in_idx += blockDim.x) {\r\n        int rel_in_y = in_idx / tile_in_w;\r\n        int rel_in_x = in_idx - rel_in_y * tile_in_w;\r\n        int in_x = rel_in_x + tile_in_x;\r\n        int in_y = rel_in_y + tile_in_y;\r\n\r\n        scalar_t v = 0.0;\r\n\r\n        if (in_x >= 0 & in_y >= 0 & in_x < p.in_w & in_y < p.in_h) {\r\n          v = input[((major_idx * p.in_h + in_y) * p.in_w + in_x) *\r\n                        p.minor_dim +\r\n                    minor_idx];\r\n        }\r\n\r\n        sx[rel_in_y][rel_in_x] = v;\r\n      }\r\n\r\n      __syncthreads();\r\n      for (int out_idx = threadIdx.x; out_idx < tile_out_h * tile_out_w;\r\n           out_idx += blockDim.x) {\r\n        int rel_out_y = out_idx / tile_out_w;\r\n        int rel_out_x = out_idx - rel_out_y * tile_out_w;\r\n        int out_x = rel_out_x + tile_out_x;\r\n        int out_y = rel_out_y + tile_out_y;\r\n\r\n        int mid_x = tile_mid_x + rel_out_x * down_x;\r\n        int mid_y = tile_mid_y + rel_out_y * down_y;\r\n        int in_x = floor_div(mid_x, up_x);\r\n        int in_y = floor_div(mid_y, up_y);\r\n        int rel_in_x = in_x - tile_in_x;\r\n        int rel_in_y = in_y - tile_in_y;\r\n        int kernel_x = (in_x + 1) * up_x - mid_x - 1;\r\n        int kernel_y = (in_y + 1) * up_y - mid_y - 1;\r\n\r\n        scalar_t v = 0.0;\r\n\r\n#pragma unroll\r\n        for (int y = 0; y < kernel_h / up_y; y++)\r\n#pragma unroll\r\n          for (int x = 0; x < kernel_w / up_x; x++)\r\n            v += sx[rel_in_y + y][rel_in_x + x] *\r\n                 sk[kernel_y + y * up_y][kernel_x + x * up_x];\r\n\r\n        if (out_x < p.out_w & out_y < p.out_h) {\r\n          out[((major_idx * p.out_h + out_y) * p.out_w + out_x) * p.minor_dim +\r\n              minor_idx] = v;\r\n        }\r\n      }\r\n    }\r\n  }\r\n}\r\n\r\ntorch::Tensor upfirdn2d_op(const torch::Tensor &input,\r\n                           const torch::Tensor &kernel, int up_x, int up_y,\r\n                           int down_x, int down_y, int pad_x0, int pad_x1,\r\n                           int pad_y0, int pad_y1) {\r\n  int curDevice = -1;\r\n  cudaGetDevice(&curDevice);\r\n  cudaStream_t stream = at::cuda::getCurrentCUDAStream(curDevice);\r\n\r\n  UpFirDn2DKernelParams p;\r\n\r\n  auto x = input.contiguous();\r\n  auto k = kernel.contiguous();\r\n\r\n  p.major_dim = x.size(0);\r\n  p.in_h = x.size(1);\r\n  p.in_w = x.size(2);\r\n  p.minor_dim = x.size(3);\r\n  p.kernel_h = k.size(0);\r\n  p.kernel_w = k.size(1);\r\n  p.up_x = up_x;\r\n  p.up_y = up_y;\r\n  p.down_x = down_x;\r\n  p.down_y = down_y;\r\n  p.pad_x0 = pad_x0;\r\n  p.pad_x1 = pad_x1;\r\n  p.pad_y0 = pad_y0;\r\n  p.pad_y1 = pad_y1;\r\n\r\n  p.out_h = (p.in_h * p.up_y + p.pad_y0 + p.pad_y1 - p.kernel_h + p.down_y) /\r\n            p.down_y;\r\n  p.out_w = (p.in_w * p.up_x + p.pad_x0 + p.pad_x1 - p.kernel_w + p.down_x) /\r\n            p.down_x;\r\n\r\n  auto out =\r\n      at::empty({p.major_dim, p.out_h, p.out_w, p.minor_dim}, x.options());\r\n\r\n  int mode = -1;\r\n\r\n  int tile_out_h = -1;\r\n  int tile_out_w = -1;\r\n\r\n  if (p.up_x == 1 && p.up_y == 1 && p.down_x == 1 && p.down_y == 1 &&\r\n      p.kernel_h <= 4 && p.kernel_w <= 4) {\r\n    mode = 1;\r\n    tile_out_h = 16;\r\n    tile_out_w = 64;\r\n  }\r\n\r\n  if (p.up_x == 1 && p.up_y == 1 && p.down_x == 1 && p.down_y == 1 &&\r\n      p.kernel_h <= 3 && p.kernel_w <= 3) {\r\n    mode = 2;\r\n    tile_out_h = 16;\r\n    tile_out_w = 64;\r\n  }\r\n\r\n  if (p.up_x == 2 && p.up_y == 2 && p.down_x == 1 && p.down_y == 1 &&\r\n      p.kernel_h <= 4 && p.kernel_w <= 4) {\r\n    mode = 3;\r\n    tile_out_h = 16;\r\n    tile_out_w = 64;\r\n  }\r\n\r\n  if (p.up_x == 2 && p.up_y == 2 && p.down_x == 1 && p.down_y == 1 &&\r\n      p.kernel_h <= 2 && p.kernel_w <= 2) {\r\n    mode = 4;\r\n    tile_out_h = 16;\r\n    tile_out_w = 64;\r\n  }\r\n\r\n  if (p.up_x == 1 && p.up_y == 1 && p.down_x == 2 && p.down_y == 2 &&\r\n      p.kernel_h <= 4 && p.kernel_w <= 4) {\r\n    mode = 5;\r\n    tile_out_h = 8;\r\n    tile_out_w = 32;\r\n  }\r\n\r\n  if (p.up_x == 1 && p.up_y == 1 && p.down_x == 2 && p.down_y == 2 &&\r\n      p.kernel_h <= 2 && p.kernel_w <= 2) {\r\n    mode = 6;\r\n    tile_out_h = 8;\r\n    tile_out_w = 32;\r\n  }\r\n\r\n  dim3 block_size;\r\n  dim3 grid_size;\r\n\r\n  if (tile_out_h > 0 && tile_out_w > 0) {\r\n    p.loop_major = (p.major_dim - 1) / 16384 + 1;\r\n    p.loop_x = 1;\r\n    block_size = dim3(32 * 8, 1, 1);\r\n    grid_size = dim3(((p.out_h - 1) / tile_out_h + 1) * p.minor_dim,\r\n                     (p.out_w - 1) / (p.loop_x * tile_out_w) + 1,\r\n                     (p.major_dim - 1) / p.loop_major + 1);\r\n  } else {\r\n    p.loop_major = (p.major_dim - 1) / 16384 + 1;\r\n    p.loop_x = 4;\r\n    block_size = dim3(4, 32, 1);\r\n    grid_size = dim3((p.out_h * p.minor_dim - 1) / block_size.x + 1,\r\n                     (p.out_w - 1) / (p.loop_x * block_size.y) + 1,\r\n                     (p.major_dim - 1) / p.loop_major + 1);\r\n  }\r\n\r\n  AT_DISPATCH_FLOATING_TYPES_AND_HALF(x.scalar_type(), \"upfirdn2d_cuda\", [&] {\r\n    switch (mode) {\r\n    case 1:\r\n      upfirdn2d_kernel<scalar_t, 1, 1, 1, 1, 4, 4, 16, 64>\r\n          <<<grid_size, block_size, 0, stream>>>(out.data_ptr<scalar_t>(),\r\n                                                 x.data_ptr<scalar_t>(),\r\n                                                 k.data_ptr<scalar_t>(), p);\r\n\r\n      break;\r\n\r\n    case 2:\r\n      upfirdn2d_kernel<scalar_t, 1, 1, 1, 1, 3, 3, 16, 64>\r\n          <<<grid_size, block_size, 0, stream>>>(out.data_ptr<scalar_t>(),\r\n                                                 x.data_ptr<scalar_t>(),\r\n                                                 k.data_ptr<scalar_t>(), p);\r\n\r\n      break;\r\n\r\n    case 3:\r\n      upfirdn2d_kernel<scalar_t, 2, 2, 1, 1, 4, 4, 16, 64>\r\n          <<<grid_size, block_size, 0, stream>>>(out.data_ptr<scalar_t>(),\r\n                                                 x.data_ptr<scalar_t>(),\r\n                                                 k.data_ptr<scalar_t>(), p);\r\n\r\n      break;\r\n\r\n    case 4:\r\n      upfirdn2d_kernel<scalar_t, 2, 2, 1, 1, 2, 2, 16, 64>\r\n          <<<grid_size, block_size, 0, stream>>>(out.data_ptr<scalar_t>(),\r\n                                                 x.data_ptr<scalar_t>(),\r\n                                                 k.data_ptr<scalar_t>(), p);\r\n\r\n      break;\r\n\r\n    case 5:\r\n      upfirdn2d_kernel<scalar_t, 1, 1, 2, 2, 4, 4, 8, 32>\r\n          <<<grid_size, block_size, 0, stream>>>(out.data_ptr<scalar_t>(),\r\n                                                 x.data_ptr<scalar_t>(),\r\n                                                 k.data_ptr<scalar_t>(), p);\r\n\r\n      break;\r\n\r\n    case 6:\r\n      upfirdn2d_kernel<scalar_t, 1, 1, 2, 2, 4, 4, 8, 32>\r\n          <<<grid_size, block_size, 0, stream>>>(out.data_ptr<scalar_t>(),\r\n                                                 x.data_ptr<scalar_t>(),\r\n                                                 k.data_ptr<scalar_t>(), p);\r\n\r\n      break;\r\n\r\n    default:\r\n      upfirdn2d_kernel_large<scalar_t><<<grid_size, block_size, 0, stream>>>(\r\n          out.data_ptr<scalar_t>(), x.data_ptr<scalar_t>(),\r\n          k.data_ptr<scalar_t>(), p);\r\n    }\r\n  });\r\n\r\n  return out;\r\n}"
  },
  {
    "path": "prepare_data.py",
    "content": "import argparse\nfrom io import BytesIO\nimport multiprocessing\nfrom functools import partial\n\nfrom PIL import Image\n\nimport lmdb\nfrom tqdm import tqdm\nfrom torchvision import datasets\nfrom torchvision.transforms import functional as trans_fn\n\n# ImageFile.LOAD_TRUNCATED_IMAGES = True\n\n\ndef resize_and_convert(img, size, resample, quality=100):\n    img = trans_fn.resize(img, size, resample)\n    img = trans_fn.center_crop(img, size)\n    buffer = BytesIO()\n    img.save(buffer, format=\"jpeg\", quality=quality)\n    val = buffer.getvalue()\n\n    return val\n\n\ndef resize_multiple(img, sizes=(128, 256, 512, 1024), resample=Image.LANCZOS, quality=100):\n    imgs = []\n\n    for size in sizes:\n        imgs.append(resize_and_convert(img, size, resample, quality))\n\n    return imgs\n\n\ndef resize_worker(img_file, sizes, resample):\n    i, file = img_file\n    try:\n        img = Image.open(file)\n        img = img.convert(\"RGB\")\n    except:\n        print(file, \"truncated\")\n    out = resize_multiple(img, sizes=sizes, resample=resample)\n\n    return i, out\n\n\ndef prepare(env, dataset, n_worker, sizes=(128, 256, 512, 1024), resample=Image.LANCZOS):\n    resize_fn = partial(resize_worker, sizes=sizes, resample=resample)\n\n    files = sorted(dataset.imgs, key=lambda x: x[0])\n    files = [(i, file) for i, (file, label) in enumerate(files)]\n    total = 0\n\n    with multiprocessing.Pool(n_worker) as pool:\n        for i, imgs in tqdm(pool.imap_unordered(resize_fn, files)):\n            for size, img in zip(sizes, imgs):\n                key = f\"{size}-{str(i).zfill(5)}\".encode(\"utf-8\")\n\n                with env.begin(write=True) as txn:\n                    txn.put(key, img)\n\n            total += 1\n\n        with env.begin(write=True) as txn:\n            txn.put(\"length\".encode(\"utf-8\"), str(total).encode(\"utf-8\"))\n\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser()\n    parser.add_argument(\"--out\", type=str)\n    parser.add_argument(\"--size\", type=str, default=\"128,256,512,1024\")\n    parser.add_argument(\"--n_worker\", type=int, default=8)\n    parser.add_argument(\"--resample\", type=str, default=\"bilinear\")\n    parser.add_argument(\"path\", type=str)\n\n    args = parser.parse_args()\n\n    resample_map = {\"lanczos\": Image.LANCZOS, \"bilinear\": Image.BILINEAR}\n    resample = resample_map[args.resample]\n\n    sizes = [int(s.strip()) for s in args.size.split(\",\")]\n\n    print(f\"Make dataset of image sizes:\", \", \".join(str(s) for s in sizes))\n\n    imgset = datasets.ImageFolder(args.path)\n\n    with lmdb.open(args.out, map_size=1024 ** 4, readahead=False) as env:\n        prepare(env, imgset, args.n_worker, sizes=sizes, resample=resample)\n"
  },
  {
    "path": "prepare_vae_codes.py",
    "content": "import argparse\nimport numpy as np\nimport multiprocessing\nfrom functools import partial\n\nimport lmdb\nfrom tqdm import tqdm\n\nimport torch as th\nfrom autoencoder import ConvSegNet\nfrom torchvision import datasets\nimport torchvision.transforms as transforms\n\n\ndef lmdmb_write_worker(i_code, env, size):\n    i, code = i_code.cpu().numpy()\n    key = f\"{size}-{str(i).zfill(5)}\".encode(\"utf-8\")\n    with env.begin(write=True) as txn:\n        txn.put(key, code)\n\n\ndef prepare(env, vae, loader, total, batch_size, n_worker, size=1024):\n    write_fn = partial(lmdmb_write_worker, env=env, size=size)\n\n    b = 0\n    with multiprocessing.Pool(n_worker) as pool:\n        for batch in tqdm(loader):\n            code_nums = np.arange(b * batch_size, (b + 1) * batch_size)\n\n            with th.no_grad():\n                codes = vae.module.encode(batch[0].cuda())\n\n            pool.imap_unordered(write_fn, zip(code_nums, codes))\n\n            b += 1\n\n    with env.begin(write=True) as txn:\n        txn.put(\"length\".encode(\"utf-8\"), str(total).encode(\"utf-8\"))\n\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser()\n    parser.add_argument(\"--out\", type=str)\n    parser.add_argument(\"--size\", type=int, default=1024)\n    parser.add_argument(\"--n_worker\", type=int, default=24)\n    parser.add_argument(\"--batch_size\", type=int, default=4)\n    parser.add_argument(\"--resample\", type=str, default=\"bilinear\")\n    parser.add_argument(\"data_path\", type=str)\n    parser.add_argument(\"vae_checkpoint\", type=str)\n\n    args = parser.parse_args()\n\n    print(f\"Make dataset of image size:\", args.size)\n\n    transform = transforms.Compose(\n        [\n            transforms.Resize(args.size),\n            transforms.ToTensor(),\n            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True),\n        ]\n    )\n    imgset = datasets.ImageFolder(args.data_path, transform=transform)\n    loader = th.utils.data.DataLoader(imgset, batch_size=args.batch_size, num_workers=int(args.n_worker / 2))\n    print(args.batch_size)\n    print(loader)\n\n    vae = ConvSegNet()\n    vae.load_state_dict(th.load(args.vae_checkpoint)[\"vae\"])\n    vae = th.nn.DataParallel(vae).eval().cuda()\n\n    with lmdb.open(args.out, map_size=1024 ** 4, readahead=False) as env:\n        prepare(\n            env,\n            vae,\n            loader,\n            total=len(imgset),\n            batch_size=args.batch_size,\n            n_worker=int(args.n_worker / 2),\n            size=args.size,\n        )\n"
  },
  {
    "path": "projector.py",
    "content": "import argparse\r\nimport math\r\nimport os\r\n\r\nimport torch\r\nfrom torch import optim\r\nfrom torch.nn import functional as F\r\nfrom torchvision import transforms\r\nfrom PIL import Image\r\nfrom tqdm import tqdm\r\n\r\nimport lpips\r\nfrom model import Generator\r\n\r\n\r\ndef noise_regularize(noises):\r\n    loss = 0\r\n\r\n    for noise in noises:\r\n        size = noise.shape[2]\r\n\r\n        while True:\r\n            loss = (\r\n                loss\r\n                + (noise * torch.roll(noise, shifts=1, dims=3)).mean().pow(2)\r\n                + (noise * torch.roll(noise, shifts=1, dims=2)).mean().pow(2)\r\n            )\r\n\r\n            if size <= 8:\r\n                break\r\n\r\n            noise = noise.reshape([1, 1, size // 2, 2, size // 2, 2])\r\n            noise = noise.mean([3, 5])\r\n            size //= 2\r\n\r\n    return loss\r\n\r\n\r\ndef noise_normalize_(noises):\r\n    for noise in noises:\r\n        mean = noise.mean()\r\n        std = noise.std()\r\n\r\n        noise.data.add_(-mean).div_(std)\r\n\r\n\r\ndef get_lr(t, initial_lr, rampdown=0.25, rampup=0.05):\r\n    lr_ramp = min(1, (1 - t) / rampdown)\r\n    lr_ramp = 0.5 - 0.5 * math.cos(lr_ramp * math.pi)\r\n    lr_ramp = lr_ramp * min(1, t / rampup)\r\n\r\n    return initial_lr * lr_ramp\r\n\r\n\r\ndef latent_noise(latent, strength):\r\n    noise = torch.randn_like(latent) * strength\r\n\r\n    return latent + noise\r\n\r\n\r\ndef make_image(tensor):\r\n    return (\r\n        tensor.detach()\r\n        .clamp_(min=-1, max=1)\r\n        .add(1)\r\n        .div_(2)\r\n        .mul(255)\r\n        .type(torch.uint8)\r\n        .permute(0, 2, 3, 1)\r\n        .to(\"cpu\")\r\n        .numpy()\r\n    )\r\n\r\n\r\nif __name__ == \"__main__\":\r\n    device = \"cuda\"\r\n\r\n    parser = argparse.ArgumentParser()\r\n    parser.add_argument(\"--ckpt\", type=str, required=True)\r\n    parser.add_argument(\"--size\", type=int, default=256)\r\n    parser.add_argument(\"--lr_rampup\", type=float, default=0.05)\r\n    parser.add_argument(\"--lr_rampdown\", type=float, default=0.25)\r\n    parser.add_argument(\"--lr\", type=float, default=0.1)\r\n    parser.add_argument(\"--noise\", type=float, default=0.05)\r\n    parser.add_argument(\"--noise_ramp\", type=float, default=0.75)\r\n    parser.add_argument(\"--step\", type=int, default=1000)\r\n    parser.add_argument(\"--noise_regularize\", type=float, default=1e5)\r\n    parser.add_argument(\"--mse\", type=float, default=0)\r\n    parser.add_argument(\"--w_plus\", action=\"store_true\")\r\n    parser.add_argument(\"files\", metavar=\"FILES\", nargs=\"+\")\r\n\r\n    args = parser.parse_args()\r\n\r\n    n_mean_latent = 10000\r\n\r\n    resize = min(args.size, 256)\r\n\r\n    transform = transforms.Compose(\r\n        [\r\n            transforms.Resize(resize),\r\n            transforms.CenterCrop(resize),\r\n            transforms.ToTensor(),\r\n            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),\r\n        ]\r\n    )\r\n\r\n    imgs = []\r\n\r\n    for imgfile in args.files:\r\n        img = transform(Image.open(imgfile).convert(\"RGB\"))\r\n        imgs.append(img)\r\n\r\n    imgs = torch.stack(imgs, 0).to(device)\r\n\r\n    g_ema = Generator(args.size, 512, 8)\r\n    g_ema.load_state_dict(torch.load(args.ckpt)[\"g_ema\"], strict=False)\r\n    g_ema.eval()\r\n    g_ema = g_ema.to(device)\r\n\r\n    with torch.no_grad():\r\n        noise_sample = torch.randn(n_mean_latent, 512, device=device)\r\n        latent_out = g_ema.style(noise_sample)\r\n\r\n        latent_mean = latent_out.mean(0)\r\n        latent_std = ((latent_out - latent_mean).pow(2).sum() / n_mean_latent) ** 0.5\r\n\r\n    percept = lpips.PerceptualLoss(model=\"net-lin\", net=\"vgg\", use_gpu=device.startswith(\"cuda\"))\r\n\r\n    noises = g_ema.make_noise()\r\n\r\n    latent_in = latent_mean.detach().clone().unsqueeze(0).repeat(2, 1)\r\n\r\n    if args.w_plus:\r\n        latent_in = latent_in.unsqueeze(1).repeat(1, g_ema.n_latent, 1)\r\n\r\n    latent_in.requires_grad = True\r\n\r\n    for noise in noises:\r\n        noise.requires_grad = True\r\n\r\n    optimizer = optim.Adam([latent_in] + noises, lr=args.lr)\r\n\r\n    pbar = tqdm(range(args.step))\r\n    latent_path = []\r\n\r\n    for i in pbar:\r\n        t = i / args.step\r\n        lr = get_lr(t, args.lr)\r\n        optimizer.param_groups[0][\"lr\"] = lr\r\n        noise_strength = latent_std * args.noise * max(0, 1 - t / args.noise_ramp) ** 2\r\n        latent_n = latent_noise(latent_in, noise_strength.item())\r\n\r\n        img_gen, _ = g_ema([latent_n], input_is_latent=True, noise=noises)\r\n\r\n        batch, channel, height, width = img_gen.shape\r\n\r\n        if height > 256:\r\n            factor = height // 256\r\n\r\n            img_gen = img_gen.reshape(batch, channel, height // factor, factor, width // factor, factor)\r\n            img_gen = img_gen.mean([3, 5])\r\n\r\n        p_loss = percept(img_gen, imgs).sum()\r\n        n_loss = noise_regularize(noises)\r\n        mse_loss = F.mse_loss(img_gen, imgs)\r\n\r\n        loss = p_loss + args.noise_regularize * n_loss + args.mse * mse_loss\r\n\r\n        optimizer.zero_grad()\r\n        loss.backward()\r\n        optimizer.step()\r\n\r\n        noise_normalize_(noises)\r\n\r\n        if (i + 1) % 100 == 0:\r\n            latent_path.append(latent_in.detach().clone())\r\n\r\n        pbar.set_description(\r\n            (\r\n                f\"perceptual: {p_loss.item():.4f}; noise regularize: {n_loss.item():.4f};\"\r\n                f\" mse: {mse_loss.item():.4f}; lr: {lr:.4f}\"\r\n            )\r\n        )\r\n\r\n    result_file = {\"noises\": noises}\r\n\r\n    img_gen, _ = g_ema([latent_path[-1]], input_is_latent=True, noise=noises)\r\n\r\n    filename = os.path.splitext(os.path.basename(args.files[0]))[0] + \".pt\"\r\n\r\n    img_ar = make_image(img_gen)\r\n\r\n    for i, input_name in enumerate(args.files):\r\n        result_file[input_name] = {\"img\": img_gen[i], \"latent\": latent_in[i]}\r\n        img_name = os.path.splitext(os.path.basename(input_name))[0] + \"-project.png\"\r\n        pil_img = Image.fromarray(img_ar[i])\r\n        pil_img.save(img_name)\r\n\r\n    torch.save(result_file, filename)\r\n"
  },
  {
    "path": "render.py",
    "content": "import queue\nfrom threading import Thread\n\nimport ffmpeg\nimport numpy as np\nimport PIL.Image\nimport torch as th\nfrom tqdm import tqdm\n\nth.set_grad_enabled(False)\nth.backends.cudnn.benchmark = True\n\n\ndef render(\n    generator,\n    latents,\n    noise,\n    offset,\n    duration,\n    batch_size,\n    out_size,\n    output_file,\n    audio_file=None,\n    truncation=1.0,\n    bends=[],\n    rewrites={},\n    randomize_noise=False,\n    ffmpeg_preset=\"slow\",\n):\n    split_queue = queue.Queue()\n    render_queue = queue.Queue()\n\n    # postprocesses batched torch tensors to individual RGB numpy arrays\n    def split_batches(jobs_in, jobs_out):\n        while True:\n            try:\n                imgs = jobs_in.get(timeout=5)\n            except queue.Empty:\n                return\n            imgs = (imgs.clamp_(-1, 1) + 1) * 127.5\n            imgs = imgs.permute(0, 2, 3, 1)\n            for img in imgs:\n                jobs_out.put(img.cpu().numpy().astype(np.uint8))\n            jobs_in.task_done()\n\n    # start background ffmpeg process that listens on stdin for frame data\n    if out_size == 512:\n        output_size = \"512x512\"\n    elif out_size == 1024:\n        output_size = \"1024x1024\"\n    elif out_size == 1920:\n        output_size = \"1920x1080\"\n    elif out_size == 1080:\n        output_size = \"1080x1920\"\n    else:\n        raise Exception(\"The only output sizes currently supported are: 512, 1024, 1080, or 1920\")\n\n    if audio_file is not None:\n        audio = ffmpeg.input(audio_file, ss=offset, t=duration, guess_layout_max=0)\n        video = (\n            ffmpeg.input(\"pipe:\", format=\"rawvideo\", pix_fmt=\"rgb24\", framerate=len(latents) / duration, s=output_size)\n            .output(\n                audio,\n                output_file,\n                framerate=len(latents) / duration,\n                vcodec=\"libx264\",\n                pix_fmt=\"yuv420p\",\n                preset=ffmpeg_preset,\n                audio_bitrate=\"320K\",\n                ac=2,\n                v=\"warning\",\n            )\n            .global_args(\"-hide_banner\")\n            .overwrite_output()\n            .run_async(pipe_stdin=True)\n        )\n    else:\n        video = (\n            ffmpeg.input(\"pipe:\", format=\"rawvideo\", pix_fmt=\"rgb24\", framerate=len(latents) / duration, s=output_size)\n            .output(\n                output_file,\n                framerate=len(latents) / duration,\n                vcodec=\"libx264\",\n                pix_fmt=\"yuv420p\",\n                preset=ffmpeg_preset,\n                v=\"warning\",\n            )\n            .global_args(\"-hide_banner\")\n            .overwrite_output()\n            .run_async(pipe_stdin=True)\n        )\n\n    # writes numpy frames to ffmpeg stdin as raw rgb24 bytes\n    def make_video(jobs_in):\n        w, h = [int(dim) for dim in output_size.split(\"x\")]\n        for _ in tqdm(range(len(latents)), position=0, leave=True, ncols=80):\n            img = jobs_in.get(timeout=5)\n            if img.shape[1] == 2048:\n                img = img[:, 112:-112, :]\n                im = PIL.Image.fromarray(img)\n                img = np.array(im.resize((1920, 1080), PIL.Image.BILINEAR))\n            elif img.shape[0] == 2048:\n                img = img[112:-112, :, :]\n                im = PIL.Image.fromarray(img)\n                img = np.array(im.resize((1080, 1920), PIL.Image.BILINEAR))\n            assert (\n                img.shape[1] == w and img.shape[0] == h\n            ), f\"\"\"generator's output image size does not match specified output size: \\n\n                got: {img.shape[1]}x{img.shape[0]}\\t\\tshould be {output_size}\"\"\"\n            video.stdin.write(img.tobytes())\n            jobs_in.task_done()\n        video.stdin.close()\n        video.wait()\n\n    splitter = Thread(target=split_batches, args=(split_queue, render_queue))\n    splitter.daemon = True\n    renderer = Thread(target=make_video, args=(render_queue,))\n    renderer.daemon = True\n\n    # make all data that needs to be loaded to the GPU float, contiguous, and pinned\n    # the entire process is severly memory-transfer bound, but at least this might help a little\n    latents = latents.float().contiguous().pin_memory()\n\n    for ni, noise_scale in enumerate(noise):\n        noise[ni] = noise_scale.float().contiguous().pin_memory() if noise_scale is not None else None\n\n    param_dict = dict(generator.named_parameters())\n    original_weights = {}\n    for param, (rewrite, modulation) in rewrites.items():\n        rewrites[param] = [rewrite, modulation.float().contiguous().pin_memory()]\n        original_weights[param] = param_dict[param].copy().cpu().float().contiguous().pin_memory()\n\n    for bend in bends:\n        if \"modulation\" in bend:\n            bend[\"modulation\"] = bend[\"modulation\"].float().contiguous().pin_memory()\n\n    if not isinstance(truncation, float):\n        truncation = truncation.float().contiguous().pin_memory()\n\n    for n in range(0, len(latents), batch_size):\n        # load batches of data onto the GPU\n        latent_batch = latents[n : n + batch_size].cuda(non_blocking=True)\n\n        noise_batch = []\n        for noise_scale in noise:\n            if noise_scale is not None:\n                noise_batch.append(noise_scale[n : n + batch_size].cuda(non_blocking=True))\n            else:\n                noise_batch.append(None)\n\n        bend_batch = []\n        if bends is not None:\n            for bend in bends:\n                if \"modulation\" in bend:\n                    transform = bend[\"transform\"](bend[\"modulation\"][n : n + batch_size].cuda(non_blocking=True))\n                    bend_batch.append({\"layer\": bend[\"layer\"], \"transform\": transform})\n                else:\n                    bend_batch.append({\"layer\": bend[\"layer\"], \"transform\": bend[\"transform\"]})\n\n        for param, (rewrite, modulation) in rewrites.items():\n            transform = rewrite(modulation[n : n + batch_size])\n            rewritten_weight = transform(original_weights[param]).cuda(non_blocking=True)\n            param_attrs = param.split(\".\")\n            mod = generator\n            for attr in param_attrs[:-1]:\n                mod = getattr(mod, attr)\n            setattr(mod, param_attrs[-1], th.nn.Parameter(rewritten_weight))\n\n        if not isinstance(truncation, float):\n            truncation_batch = truncation[n : n + batch_size].cuda(non_blocking=True)\n        else:\n            truncation_batch = truncation\n\n        # forward through the generator\n        outputs, _ = generator(\n            styles=latent_batch,\n            noise=noise_batch,\n            truncation=truncation_batch,\n            transform_dict_list=bend_batch,\n            randomize_noise=randomize_noise,\n            input_is_latent=True,\n        )\n\n        # send output to be split into frames and rendered one by one\n        split_queue.put(outputs)\n\n        if n == 0:\n            splitter.start()\n            renderer.start()\n\n    splitter.join()\n    renderer.join()\n\n\ndef write_video(arr, output_file, fps):\n    print(f\"writing {arr.shape[0]} frames...\")\n\n    output_size = \"x\".join(reversed([str(s) for s in arr.shape[1:-1]]))\n\n    ffmpeg_proc = (\n        ffmpeg.input(\"pipe:\", format=\"rawvideo\", pix_fmt=\"rgb24\", framerate=fps, s=output_size)\n        .output(output_file, framerate=fps, vcodec=\"libx264\", preset=\"slow\", v=\"warning\")\n        .global_args(\"-benchmark\", \"-stats\", \"-hide_banner\")\n        .overwrite_output()\n        .run_async(pipe_stdin=True)\n    )\n\n    for frame in arr:\n        ffmpeg_proc.stdin.write(frame.astype(np.uint8).tobytes())\n\n    ffmpeg_proc.stdin.close()\n    ffmpeg_proc.wait()\n"
  },
  {
    "path": "requirements.txt",
    "content": "torch\ntorchvision\nnumpy\nlibrosa\ncython\nmadmom\ntqdm\nkornia\nmatplotlib\nffmpeg-python\njoblib"
  },
  {
    "path": "select_latents.py",
    "content": "import gc, math\nimport argparse\nimport tkinter as tk\nimport numpy as np\nfrom PIL import Image, ImageTk\nimport torch as th\nimport torch.nn.functional as F\nimport torchvision\nfrom models.stylegan2 import Generator as G_style2\nimport tkinter as tk\n\n# --- classes ---\ntry:\n    from Tkinter import Canvas, Frame\n    from ttk import Scrollbar\n\n    from Tkconstants import *\nexcept ImportError:\n    from tkinter import Canvas, Frame\n    from tkinter.ttk import Scrollbar\n\n    from tkinter.constants import *\n\nimport platform\n\nOS = platform.system()\n\n\nclass HoverButton(tk.Button):\n    def __init__(self, master, **kw):\n        tk.Button.__init__(self, master=master, **kw)\n        self.defaultBackground = self[\"background\"]\n        self.bind(\"<Enter>\", self.on_enter)\n        self.bind(\"<Leave>\", self.on_leave)\n\n    def on_enter(self, e):\n        self[\"background\"] = self[\"activebackground\"]\n\n    def on_leave(self, e):\n        self[\"background\"] = self.defaultBackground\n\n\nclass InvisibleScrollbar(Scrollbar):\n    def set(self, lo, hi):\n        self.tk.call(\"grid\", \"remove\", self)\n        Scrollbar.set(self, lo, hi)\n\n\nclass Mousewheel_Support(object):\n    # implemetation of singleton pattern\n    _instance = None\n\n    def __new__(cls, *args, **kwargs):\n        if not cls._instance:\n            cls._instance = object.__new__(cls)\n        return cls._instance\n\n    def __init__(self, root, horizontal_factor=1, vertical_factor=1):\n        self._active_area = None\n        if isinstance(horizontal_factor, int):\n            self.horizontal_factor = horizontal_factor\n        else:\n            raise Exception(\"Vertical factor must be an integer.\")\n        if isinstance(vertical_factor, int):\n            self.vertical_factor = vertical_factor\n        else:\n            raise Exception(\"Horizontal factor must be an integer.\")\n        if OS == \"Linux\":\n            root.bind_all(\"<4>\", self._on_mousewheel, add=\"+\")\n            root.bind_all(\"<5>\", self._on_mousewheel, add=\"+\")\n        else:\n            # Windows and MacOS\n            root.bind_all(\"<MouseWheel>\", self._on_mousewheel, add=\"+\")\n\n    def _on_mousewheel(self, event):\n        if self._active_area:\n            self._active_area.onMouseWheel(event)\n\n    def _mousewheel_bind(self, widget):\n        self._active_area = widget\n\n    def _mousewheel_unbind(self):\n        self._active_area = None\n\n    def add_support_to(\n        self, widget=None, xscrollbar=None, yscrollbar=None, what=\"units\", horizontal_factor=None, vertical_factor=None\n    ):\n        if xscrollbar is None and yscrollbar is None:\n            return\n        if xscrollbar is not None:\n            horizontal_factor = horizontal_factor or self.horizontal_factor\n            xscrollbar.onMouseWheel = self._make_mouse_wheel_handler(widget, \"x\", self.horizontal_factor, what)\n            xscrollbar.bind(\"<Enter>\", lambda event, scrollbar=xscrollbar: self._mousewheel_bind(scrollbar))\n            xscrollbar.bind(\"<Leave>\", lambda event: self._mousewheel_unbind())\n        if yscrollbar is not None:\n            vertical_factor = vertical_factor or self.vertical_factor\n            yscrollbar.onMouseWheel = self._make_mouse_wheel_handler(widget, \"y\", self.vertical_factor, what)\n            yscrollbar.bind(\"<Enter>\", lambda event, scrollbar=yscrollbar: self._mousewheel_bind(scrollbar))\n            yscrollbar.bind(\"<Leave>\", lambda event: self._mousewheel_unbind())\n        main_scrollbar = yscrollbar if yscrollbar is not None else xscrollbar\n        if widget is not None:\n            if isinstance(widget, list) or isinstance(widget, tuple):\n                list_of_widgets = widget\n                for widget in list_of_widgets:\n                    widget.bind(\"<Enter>\", lambda event: self._mousewheel_bind(widget))\n                    widget.bind(\"<Leave>\", lambda event: self._mousewheel_unbind())\n                    widget.onMouseWheel = main_scrollbar.onMouseWheel\n            else:\n                widget.bind(\"<Enter>\", lambda event: self._mousewheel_bind(widget))\n                widget.bind(\"<Leave>\", lambda event: self._mousewheel_unbind())\n                widget.onMouseWheel = main_scrollbar.onMouseWheel\n\n    @staticmethod\n    def _make_mouse_wheel_handler(widget, orient, factor=1 / 120, what=\"units\"):\n        view_command = getattr(widget, orient + \"view\")\n        if OS == \"Linux\":\n\n            def onMouseWheel(event):\n                if event.num == 4:\n                    view_command(\"scroll\", (-1) * factor, what)\n                elif event.num == 5:\n                    view_command(\"scroll\", factor, what)\n\n        elif OS == \"Windows\":\n\n            def onMouseWheel(event):\n                view_command(\"scroll\", (-1) * int((event.delta / 120) * factor), what)\n\n        elif OS == \"Darwin\":\n\n            def onMouseWheel(event):\n                view_command(\"scroll\", event.delta, what)\n\n        return onMouseWheel\n\n\nclass Scrolling_Area(Frame, object):\n    def __init__(\n        self,\n        master,\n        width=None,\n        anchor=N,\n        height=None,\n        mousewheel_speed=2,\n        scroll_horizontally=True,\n        xscrollbar=None,\n        scroll_vertically=True,\n        yscrollbar=None,\n        background=\"black\",\n        inner_frame=Frame,\n        **kw,\n    ):\n        Frame.__init__(self, master, class_=\"Scrolling_Area\", background=background)\n        self.grid_columnconfigure(0, weight=1)\n        self.grid_rowconfigure(0, weight=1)\n        self._width = width\n        self._height = height\n        self.canvas = Canvas(self, background=background, highlightthickness=0, width=width, height=height)\n        self.canvas.grid(row=0, column=0, sticky=N + E + W + S)\n        if scroll_vertically:\n            if yscrollbar is not None:\n                self.yscrollbar = yscrollbar\n            else:\n                self.yscrollbar = InvisibleScrollbar(self, orient=VERTICAL)\n                self.yscrollbar.grid(row=0, column=1, sticky=N + S)\n            self.canvas.configure(yscrollcommand=self.yscrollbar.set)\n            self.yscrollbar[\"command\"] = self.canvas.yview\n        else:\n            self.yscrollbar = None\n        if scroll_horizontally:\n            if xscrollbar is not None:\n                self.xscrollbar = xscrollbar\n            else:\n                self.xscrollbar = InvisibleScrollbar(self, orient=HORIZONTAL)\n                self.xscrollbar.grid(row=1, column=0, sticky=E + W)\n            self.canvas.configure(xscrollcommand=self.xscrollbar.set)\n            self.xscrollbar[\"command\"] = self.canvas.xview\n        else:\n            self.xscrollbar = None\n        self.rowconfigure(0, weight=1)\n        self.columnconfigure(0, weight=1)\n        self.innerframe = inner_frame(self.canvas, **kw)\n        self.innerframe.pack(anchor=anchor)\n        self.canvas.create_window(0, 0, window=self.innerframe, anchor=\"nw\", tags=\"inner_frame\")\n        self.canvas.bind(\"<Configure>\", self._on_canvas_configure)\n        Mousewheel_Support(self).add_support_to(self.canvas, xscrollbar=self.xscrollbar, yscrollbar=self.yscrollbar)\n\n    @property\n    def width(self):\n        return self.canvas.winfo_width()\n\n    @width.setter\n    def width(self, width):\n        self.canvas.configure(width=width)\n\n    @property\n    def height(self):\n        return self.canvas.winfo_height()\n\n    @height.setter\n    def height(self, height):\n        self.canvas.configure(height=height)\n\n    def set_size(self, width, height):\n        self.canvas.configure(width=width, height=height)\n\n    def _on_canvas_configure(self, event):\n        width = max(self.innerframe.winfo_reqwidth(), event.width)\n        height = max(self.innerframe.winfo_reqheight(), event.height)\n        self.canvas.configure(scrollregion=\"0 0 %s %s\" % (width, height))\n        self.canvas.itemconfigure(\"inner_frame\", width=width, height=height)\n\n    def update_viewport(self):\n        self.update()\n        window_width = self.innerframe.winfo_reqwidth()\n        window_height = self.innerframe.winfo_reqheight()\n        if self._width is None:\n            canvas_width = window_width\n        else:\n            canvas_width = min(self._width, window_width)\n        if self._height is None:\n            canvas_height = window_height\n        else:\n            canvas_height = min(self._height, window_height)\n        self.canvas.configure(\n            scrollregion=\"0 0 %s %s\" % (window_width, window_height), width=self._width, height=self._height\n        )\n        self.canvas.itemconfigure(\"inner_frame\", width=window_width, height=window_height)\n\n\nth.set_grad_enabled(False)\nth.backends.cudnn.benchmark = True\n\nparser = argparse.ArgumentParser()\n\nparser.add_argument(\"--ckpt\", type=str)\nparser.add_argument(\"--res\", type=int, default=1024)\nparser.add_argument(\"--output_dir\", type=str, default=\"./workspace/\")\nparser.add_argument(\"--truncation\", type=float, default=1.5)\nparser.add_argument(\"--noconst\", action=\"store_false\")\n\nargs = parser.parse_args()\n\nname = args.ckpt.split(\"/\")[-1].split(\".\")[0]\nGENERATOR = (\n    G_style2(size=args.res, style_dim=512, n_mlp=8, checkpoint=args.ckpt, output_size=1024, constant_input=args.noconst)\n    .eval()\n    .cuda()\n)\n# GENERATOR = G_style(checkpoint=args.ckpt, output_size=1024).eval().cuda()\n# GENERATOR = th.nn.DataParallel(GENERATOR)\n\nIMAGES_PER_ROW = 4\nIMSIZE = (1920 - 240) // IMAGES_PER_ROW\n\nALL_LATENTS = []\nDROP_IDXS = []\nINTRO_IDXS = []\nIMAGES = []\n\n\ndef generate_images(n):\n    imgs = []\n    for _ in range(n // 8):\n        random_latents = th.randn(8, 512).cuda()\n        mapped_latents = GENERATOR(random_latents, noise=None, truncation=args.truncation, map_latents=True)\n        for latent in mapped_latents:\n            ALL_LATENTS.append(latent[None, ...].cpu().numpy())\n        batch, _ = GENERATOR(\n            styles=mapped_latents,\n            noise=None,\n            truncation=args.truncation,\n            transform_dict_list=[],\n            randomize_noise=True,\n            input_is_latent=True,\n        )\n        imgs.append(batch)\n    imgs = th.cat(imgs)[:n]\n    imgs = F.interpolate(imgs, IMSIZE, mode=\"bilinear\", align_corners=False)\n    imgs = (imgs.clamp_(-1, 1) + 1) * 127.5\n    imgs = imgs.permute(0, 2, 3, 1)\n    imgs_np = imgs.cpu().numpy().astype(np.uint8)\n    del imgs\n    gc.collect()\n    th.cuda.empty_cache()\n    return imgs_np\n\n\nroot = tk.Tk()\nroot.title(name)\n\nimgrid = Scrolling_Area(root, bg=\"black\", width=1680, height=1080)\nimgrid.pack(side=\"left\", expand=True, fill=\"both\")\n\npanel = tk.Frame(root, relief=\"flat\", bg=\"black\")\npanel.pack(side=\"right\", expand=True, fill=\"both\")\n\n\ndef render_latents(latents):\n    imgs = []\n    for i in range(latents.shape[0] // 8 + 1):\n        if len(latents[8 * i : 8 * (i + 1)]) < 1:\n            continue\n        batch, _ = GENERATOR(\n            styles=latents[8 * i : 8 * (i + 1)].cuda(),\n            noise=None,\n            truncation=args.truncation,\n            transform_dict_list=[],\n            randomize_noise=True,\n            input_is_latent=True,\n        )\n        imgs.append(batch)\n    imgs = th.cat(imgs)\n    imgs = (imgs.clamp_(-1, 1) + 1) / 2\n    return imgs\n\n\ndef save():\n    intro_latents = np.concatenate(ALL_LATENTS)[INTRO_IDXS]\n    torchvision.utils.save_image(\n        render_latents(th.from_numpy(intro_latents)),\n        f\"{args.output_dir}/{name}_intro_latents.jpg\",\n        nrow=int(round(math.sqrt(intro_latents.shape[0]) * 4 / 3)),\n        padding=0,\n        normalize=False,\n    )\n    np.save(f\"{args.output_dir}/{name}_intro_latents.npy\", intro_latents)\n\n    drop_latents = np.concatenate(ALL_LATENTS)[DROP_IDXS]\n    torchvision.utils.save_image(\n        render_latents(th.from_numpy(drop_latents)),\n        f\"{args.output_dir}/{name}_drop_latents.jpg\",\n        nrow=int(round(math.sqrt(drop_latents.shape[0]) * 4 / 3)),\n        padding=0,\n        normalize=False,\n    )\n    np.save(f\"{args.output_dir}/{name}_drop_latents.npy\", drop_latents)\n\n\ntk.Label(panel, text=\"latents\", height=3, bg=\"black\", fg=\"white\").pack(side=\"top\")\n\nbut = HoverButton(\n    panel,\n    text=\"Save\",\n    command=save,\n    height=3,\n    width=8,\n    bg=\"black\",\n    fg=\"white\",\n    activebackground=\"#333333\",\n    activeforeground=\"white\",\n    relief=\"flat\",\n    highlightbackground=\"#333333\",\n)\nbut.pack(side=\"bottom\")\n\nintro = tk.LabelFrame(panel, text=\"intro\", width=240, height=490, bg=\"black\", fg=\"white\", relief=\"flat\")\nintro.pack(side=\"top\", fill=\"both\")\n\ndrop = tk.LabelFrame(panel, text=\"drop\", width=240, height=490, bg=\"black\", fg=\"white\", relief=\"flat\")\ndrop.pack(side=\"bottom\", fill=\"both\")\n\nintrogrid = Scrolling_Area(intro, width=240, height=490, bg=\"black\")\nintrogrid.pack(side=\"top\", fill=\"both\")\n\ndropgrid = Scrolling_Area(drop, width=240, height=490, bg=\"black\")\ndropgrid.pack(side=\"bottom\", fill=\"both\")\n\nim_num = 0\nintro_im_num = 0\ndrop_im_num = 0\n\n\ndef add_intro(label):\n    global intro_im_num\n    img_id = int(label.__str__().split(\".\")[-1])\n\n    INTRO_IDXS.append(img_id)\n\n    img = ImageTk.PhotoImage(image=IMAGES[img_id].resize((46, 46), Image.ANTIALIAS))\n    lbl = tk.Label(introgrid.innerframe, image=img, borderwidth=0, highlightthickness=0)\n    lbl.image = img  # this line need to prevent gc\n    lbl.grid(row=math.floor(intro_im_num / 5), column=intro_im_num % 5)\n    lbl.bind(\"<Button-1>\", lambda event, l=lbl: remove_intro(l))\n\n    intro_im_num += 1\n    introgrid.update_viewport()\n\n\ndef remove_intro(label):\n    remove_idx = list(reversed(introgrid.innerframe.grid_slaves())).index(label)\n    label.grid_remove()\n    del INTRO_IDXS[remove_idx]\n    global intro_im_num\n    intro_im_num = 0\n    for im in reversed(introgrid.innerframe.grid_slaves()):\n        im.grid_remove()\n        im.grid(row=math.floor(intro_im_num / 5), column=intro_im_num % 5)\n        intro_im_num += 1\n    introgrid.update_viewport()\n\n\ndef add_drop(label):\n    global drop_im_num\n    img_id = int(label.__str__().split(\".\")[-1])\n\n    DROP_IDXS.append(img_id)\n\n    img = ImageTk.PhotoImage(image=IMAGES[img_id].resize((46, 46), Image.ANTIALIAS))\n    lbl = tk.Label(dropgrid.innerframe, image=img, borderwidth=0, highlightthickness=0)\n    lbl.image = img  # this line need to prevent gc\n    lbl.grid(row=math.floor(drop_im_num / 5), column=drop_im_num % 5)\n    lbl.bind(\"<Button-1>\", lambda event, l=lbl: remove_drop(l))\n\n    drop_im_num += 1\n    dropgrid.update_viewport()\n\n\ndef remove_drop(label):\n    remove_idx = list(reversed(dropgrid.innerframe.grid_slaves())).index(label)\n    label.grid_remove()\n    del DROP_IDXS[remove_idx]\n    global drop_im_num\n    drop_im_num = 0\n    for im in reversed(dropgrid.innerframe.grid_slaves()):\n        im.grid_remove()\n        im.grid(row=math.floor(drop_im_num / 5), column=drop_im_num % 5)\n        drop_im_num += 1\n    dropgrid.update_viewport()\n\n\ndef add_images(n):\n    global im_num, IMAGES\n\n    for im_arr in generate_images(n):\n        im = Image.fromarray(im_arr)\n        IMAGES.append(im)\n        img = ImageTk.PhotoImage(image=im)\n\n        label = tk.Label(imgrid.innerframe, image=img, name=str(im_num), borderwidth=0, highlightthickness=0)\n        label.image = img  # this line need to prevent gc\n        label.grid(row=math.floor(im_num / IMAGES_PER_ROW), column=im_num % IMAGES_PER_ROW)\n        label.bind(\"<Button-1>\", lambda event, l=label: add_intro(l))\n        label.bind(\"<Button-3>\", lambda event, l=label: add_drop(l))\n\n        im_num += 1\n\n    HoverButton(\n        imgrid.innerframe,\n        text=\"More\",\n        command=lambda n=35: add_images(n),\n        height=3,\n        width=8,\n        bg=\"black\",\n        fg=\"white\",\n        activebackground=\"#333333\",\n        activeforeground=\"white\",\n        relief=\"flat\",\n        highlightbackground=\"#333333\",\n    ).grid(\n        row=math.floor(im_num / IMAGES_PER_ROW) + 1,\n        column=math.floor(IMAGES_PER_ROW / 2 - 1),\n        columnspan=1 if math.floor(IMAGES_PER_ROW / 2 - 1) % 2 == 0 else 2,\n    )\n    imgrid.update_viewport()\n\n\nadd_images(24)\n\n\nroot.mainloop()\n"
  },
  {
    "path": "train.py",
    "content": "import argparse\nimport gc\nimport math\nimport os\nimport random\nimport sys\nimport time\n\nimport numpy as np\nimport torch as th\nimport wandb\nfrom contrastive_learner import ContrastiveLearner, RandomApply\nfrom kornia import augmentation as augs\nfrom scipy.ndimage import gaussian_filter\nfrom torch import nn\nfrom torch.nn import functional as F\nfrom torch.utils import data\nfrom torchvision import transforms, utils\nfrom tqdm import tqdm\n\nimport validation\nfrom augment import augment\nfrom dataset import MultiResolutionDataset\nfrom distributed import get_rank, reduce_loss_dict, reduce_sum, synchronize\nfrom lookahead_minimax import LookaheadMinimax\nfrom models.stylegan2 import Discriminator, Generator\n\nsys.path.insert(0, \"../lookahead_minimax\")\n\n\ndef data_sampler(dataset, shuffle, distributed):\n    if distributed:\n        return data.distributed.DistributedSampler(dataset, shuffle=shuffle)\n    if shuffle:\n        return data.RandomSampler(dataset)\n    else:\n        return data.SequentialSampler(dataset)\n\n\ndef requires_grad(model, flag=True):\n    for p in model.parameters():\n        p.requires_grad = flag\n\n\ndef accumulate(model1, model2, decay=0.5 ** (32.0 / 10_000)):\n    par1 = dict(model1.named_parameters())\n    par2 = dict(model2.named_parameters())\n    for name, param in model1.named_parameters():\n        param.data = decay * par1[name].data + (1 - decay) * par2[name].data\n\n\ndef sample_data(loader):\n    while True:\n        for batch in loader:\n            yield batch\n\n\ndef make_noise(batch_size, latent_dim, prob):\n    if prob > 0 and random.random() < prob:\n        return th.randn(2, batch_size, latent_dim, device=device).unbind(0)\n    else:\n        return [th.randn(batch_size, latent_dim, device=device)]\n\n\ndef d_logistic_loss(real_pred, fake_pred):\n    real_loss = F.softplus(-real_pred)\n    fake_loss = F.softplus(fake_pred)\n    return real_loss.mean() + fake_loss.mean()\n\n\ndef d_r1_penalty(real_img, real_pred, args):\n    (grad_real,) = th.autograd.grad(real_pred.sum(), real_img, create_graph=True)\n    r1_loss = grad_real.pow(2).view(grad_real.shape[0], -1).sum(1).mean()\n    r1_loss = r1_loss / 2.0 + 0 * real_pred[0]\n    return r1_loss\n\n\ndef g_non_saturating_loss(fake_pred):\n    return F.softplus(-fake_pred).mean()\n\n\ndef g_path_length_regularization(generator, mean_path_length, args):\n    path_batch_size = max(1, args.batch_size // args.path_batch_shrink)\n\n    noise = make_noise(path_batch_size, args.latent_size, args.mixing_prob)\n    fake_img, latents = generator(noise, return_latents=True)\n\n    img_noise = th.randn_like(fake_img) / math.sqrt(fake_img.shape[2] * fake_img.shape[3])\n    noisy_img_sum = (fake_img * img_noise).sum()\n\n    (grad,) = th.autograd.grad(noisy_img_sum, latents, create_graph=True)\n\n    path_lengths = th.sqrt(grad.pow(2).sum(2).mean(1))\n    path_mean = mean_path_length + 0.01 * (path_lengths.mean() - mean_path_length)\n    path_loss = (path_lengths - path_mean).pow(2).mean()\n    if not th.isnan(path_mean):\n        mean_path_length = path_mean.detach()\n\n    if args.path_batch_shrink:\n        path_loss += 0 * fake_img[0, 0, 0, 0]\n\n    return path_loss, mean_path_length\n\n\ndef train(args, loader, generator, discriminator, contrast_learner, g_optim, d_optim, g_ema):\n    if args.distributed:\n        g_module = generator.module\n        d_module = discriminator.module\n        if contrast_learner is not None:\n            cl_module = contrast_learner.module\n    else:\n        g_module = generator\n        d_module = discriminator\n        cl_module = contrast_learner\n\n    loader = sample_data(loader)\n    sample_z = th.randn(args.n_sample, args.latent_size, device=device)\n    mse = th.nn.MSELoss()\n    mean_path_length = th.cuda.FloatTensor([0.0])\n    ada_aug_signs = th.cuda.FloatTensor([0.0])\n    ada_aug_n = th.cuda.FloatTensor([0.0])\n    ada_aug_p = th.cuda.FloatTensor([args.augment_p if args.augment_p > 0 else 0.0])\n    ada_aug_step = th.cuda.FloatTensor([args.ada_target / args.ada_length])\n    r_t_stat = th.cuda.FloatTensor([0.0])\n    fids = []\n\n    loss_dict = {\n        \"Generator\": th.cuda.FloatTensor([0.0]),\n        \"Discriminator\": th.cuda.FloatTensor([0.0]),\n        \"Real Score\": th.cuda.FloatTensor([0.0]),\n        \"Fake Score\": th.cuda.FloatTensor([0.0]),\n        \"Contrastive\": th.cuda.FloatTensor([0.0]),\n        \"Consistency\": th.cuda.FloatTensor([0.0]),\n        \"R1 Penalty\": th.cuda.FloatTensor([0.0]),\n        \"Path Length Regularization\": th.cuda.FloatTensor([0.0]),\n        \"Augment\": th.cuda.FloatTensor([0.0]),\n        \"Rt\": th.cuda.FloatTensor([0.0]),\n    }\n\n    pbar = range(args.iter)\n    if get_rank() == 0:\n        pbar = tqdm(pbar, initial=args.start_iter, dynamic_ncols=True, smoothing=0)\n    for idx in pbar:\n        i = idx + args.start_iter\n        if i > args.iter:\n            print(\"Done!\")\n            break\n        tick_start = time.time()\n\n        for k, v in loss_dict.items():\n            loss_dict[k].mul_(0)\n\n        requires_grad(generator, False)\n        requires_grad(discriminator, True)\n\n        discriminator.zero_grad()\n        for _ in range(args.num_accumulate):\n            real_img_og = next(loader).to(device, non_blocking=True)\n\n            noise = make_noise(args.batch_size, args.latent_size, args.mixing_prob)\n            fake_img_og, _ = generator(noise)\n            if args.augment:\n                fake_img, _ = augment(fake_img_og, ada_aug_p)\n                real_img, _ = augment(real_img_og, ada_aug_p)\n            else:\n                fake_img = fake_img_og\n                real_img = real_img_og\n\n            fake_pred = discriminator(fake_img)\n            real_pred = discriminator(real_img)\n            logistic_loss = d_logistic_loss(real_pred, fake_pred)\n            loss_dict[\"Discriminator\"] += logistic_loss.detach()\n            loss_dict[\"Real Score\"] += real_pred.mean().detach()\n            loss_dict[\"Fake Score\"] += fake_pred.mean().detach()\n            d_loss = logistic_loss\n\n            if args.contrastive > 0:\n                contrast_learner(fake_img_og, fake_img, accumulate=True)\n                contrast_learner(real_img_og, real_img, accumulate=True)\n                contrast_loss = cl_module.calculate_loss()\n                loss_dict[\"Contrastive\"] += contrast_loss.detach()\n                d_loss += args.contrastive * contrast_loss\n\n            if args.balanced_consistency > 0:\n                consistency_loss = mse(real_pred, discriminator(real_img_og)) + mse(\n                    fake_pred, discriminator(fake_img_og)\n                )\n                loss_dict[\"Consistency\"] += consistency_loss.detach()\n                d_loss += args.balanced_consistency * consistency_loss\n\n            d_loss /= args.num_accumulate\n            d_loss.backward()\n        d_optim.step()\n\n        if args.r1 > 0 and i % args.d_reg_every == 0:\n            discriminator.zero_grad()\n            for _ in range(args.num_accumulate):\n                real_img = next(loader).to(device, non_blocking=True)\n                real_img.requires_grad = True\n                real_pred = discriminator(real_img)\n                r1_loss = d_r1_penalty(real_img, real_pred, args)\n                loss_dict[\"R1 Penalty\"] += r1_loss.detach().squeeze()\n                r1_loss = args.r1 * args.d_reg_every * r1_loss / args.num_accumulate\n                r1_loss.backward()\n            d_optim.step()\n\n        if args.augment and args.augment_p == 0:\n            ada_aug_signs += th.sign(real_pred).sum().item()\n            ada_aug_n += real_pred.shape[0]\n            ada_aug_signs, ada_aug_n = reduce_sum(ada_aug_signs), reduce_sum(ada_aug_n)\n\n            if ada_aug_n > 255:\n                r_t_stat = ada_aug_signs / ada_aug_n\n                loss_dict[\"Rt\"] += r_t_stat\n                if r_t_stat > args.ada_target:\n                    sign = 1\n                else:\n                    sign = -1\n\n                ada_aug_p += sign * ada_aug_step * ada_aug_n\n                ada_aug_p = th.clamp(ada_aug_p, 0, 1)\n                ada_aug_signs.mul_(0)\n                ada_aug_n.mul_(0)\n                loss_dict[\"Augment\"] += ada_aug_p\n\n        requires_grad(generator, True)\n        requires_grad(discriminator, False)\n\n        generator.zero_grad()\n        for _ in range(args.num_accumulate):\n            noise = make_noise(args.batch_size, args.latent_size, args.mixing_prob)\n            fake_img, _ = generator(noise)\n            if args.augment:\n                fake_img, _ = augment(fake_img, ada_aug_p)\n            fake_pred = discriminator(fake_img)\n            g_loss = g_non_saturating_loss(fake_pred)\n            loss_dict[\"Generator\"] += g_loss.detach()\n            g_loss /= args.num_accumulate\n            g_loss.backward()\n        g_optim.step()\n\n        if args.path_regularize > 0 and i % args.g_reg_every == 0:\n            generator.zero_grad()\n            for _ in range(args.num_accumulate):\n                path_loss, mean_path_length = g_path_length_regularization(generator, mean_path_length, args)\n                loss_dict[\"Path Length Regularization\"] += path_loss.detach()\n                path_loss = args.path_regularize * args.g_reg_every * path_loss / args.num_accumulate\n                path_loss.backward()\n            g_optim.step()\n\n        accumulate(g_ema, g_module)\n\n        loss_reduced = reduce_loss_dict(loss_dict)\n        log_dict = {k: v.mean().item() / args.num_accumulate for k, v in loss_reduced.items() if v != 0}\n        log_dict[\"Tick Length\"] = time.time() - tick_start\n\n        if get_rank() == 0:\n            with th.no_grad():\n                if args.log_spec_norm:\n                    G_norms = []\n                    for name, spec_norm in g_module.named_buffers():\n                        if \"spectral_norm\" in name:\n                            G_norms.append(spec_norm.cpu().numpy())\n                    G_norms = np.array(G_norms)\n                    D_norms = []\n                    for name, spec_norm in d_module.named_buffers():\n                        if \"spectral_norm\" in name:\n                            D_norms.append(spec_norm.cpu().numpy())\n                    D_norms = np.array(D_norms)\n                    log_dict[f\"Spectral Norms/G min spectral norm\"] = np.log(G_norms).min()\n                    log_dict[f\"Spectral Norms/G mean spectral norm\"] = np.log(G_norms).mean()\n                    log_dict[f\"Spectral Norms/G max spectral norm\"] = np.log(G_norms).max()\n                    log_dict[f\"Spectral Norms/D min spectral norm\"] = np.log(D_norms).min()\n                    log_dict[f\"Spectral Norms/D mean spectral norm\"] = np.log(D_norms).mean()\n                    log_dict[f\"Spectral Norms/D max spectral norm\"] = np.log(D_norms).max()\n\n                if args.img_every != -1 and i % args.img_every == 0:\n                    g_ema.eval()\n                    sample = []\n                    for sub in range(0, len(sample_z), args.batch_size):\n                        subsample, _ = g_ema([sample_z[sub : sub + args.batch_size]])\n                        sample.append(subsample.detach().cpu())\n                    sample = th.cat(sample).detach()\n                    grid = utils.make_grid(sample, nrow=10, normalize=True, range=(-1, 1))\n                    log_dict[\"Generated Images EMA\"] = [wandb.Image(grid, caption=f\"Step {i}\")]\n\n                if args.eval_every != -1 and i % args.eval_every == 0:\n                    fid_dict = validation.fid(\n                        g_ema, args.val_batch_size, args.fid_n_sample, args.fid_truncation, args.name\n                    )\n\n                    fid = fid_dict[\"FID\"]\n                    fids.append(fid)\n                    density = fid_dict[\"Density\"]\n                    coverage = fid_dict[\"Coverage\"]\n\n                    ppl = validation.ppl(\n                        g_ema, args.val_batch_size, args.ppl_n_sample, args.ppl_space, args.ppl_crop, args.latent_size,\n                    )\n\n                    log_dict[\"Evaluation/FID\"] = fid\n                    log_dict[\"Sweep/FID_smooth\"] = gaussian_filter(np.array(fids), [5])[-1]\n                    log_dict[\"Evaluation/Density\"] = density\n                    log_dict[\"Evaluation/Coverage\"] = coverage\n                    log_dict[\"Evaluation/PPL\"] = ppl\n\n                wandb.log(log_dict)\n\n                if args.eval_every != -1:\n                    description = (\n                        f\"FID: {fid:.4f}   PPL: {ppl:.4f}   Dens: {density:.4f}   Cov: {coverage:.4f}   \"\n                        + f\"G: {log_dict['Generator']:.4f}   D: {log_dict['Discriminator']:.4f}\"\n                    )\n                else:\n                    description = f\"G: {log_dict['Generator']:.4f}   D: {log_dict['Discriminator']:.4f}\"\n                if \"Augment\" in log_dict:\n                    description += f\"   Aug: {log_dict['Augment']:.4f}\"  #   Rt: {log_dict['Rt']:.4f}\"\n                if \"R1 Penalty\" in log_dict:\n                    description += f\"   R1: {log_dict['R1 Penalty']:.4f}\"\n                if \"Path Length Regularization\" in log_dict:\n                    description += f\"   Path: {log_dict['Path Length Regularization']:.4f}\"\n                pbar.set_description(description)\n\n                if i % args.checkpoint_every == 0:\n                    check_name = \"-\".join(\n                        [\n                            args.name,\n                            args.wbname,\n                            wandb.run.dir.split(\"/\")[-1].split(\"-\")[-1],\n                            # str(int(fid)),\n                            str(args.size),\n                            str(i).zfill(6),\n                        ]\n                    )\n                    th.save(\n                        {\n                            \"g\": g_module.state_dict(),\n                            \"d\": d_module.state_dict(),\n                            # \"cl\": cl_module.state_dict(),\n                            \"g_ema\": g_ema.state_dict(),\n                            \"g_optim\": g_optim.state_dict(),\n                            \"d_optim\": d_optim.state_dict(),\n                        },\n                        f\"/home/hans/modelzoo/maua-sg2/{check_name}.pt\",\n                    )\n\n        if args.profile_mem:\n            gpu_profile(frame=sys._getframe(), event=\"line\", arg=None)\n\n\nif __name__ == \"__main__\":\n    device = \"cuda\"\n\n    parser = argparse.ArgumentParser()\n\n    parser.add_argument(\"--wbname\", type=str, required=True)\n    parser.add_argument(\"--wbproj\", type=str, required=True)\n    parser.add_argument(\"--wbgroup\", type=str, default=None)\n\n    # data options\n    parser.add_argument(\"--path\", type=str, required=True)\n    parser.add_argument(\"--vflip\", type=bool, default=False)\n    parser.add_argument(\"--hflip\", type=bool, default=True)\n\n    # training options\n    parser.add_argument(\"--batch_size\", type=int, default=12)\n    parser.add_argument(\"--num_accumulate\", type=int, default=1)\n\n    parser.add_argument(\"--checkpoint\", type=str, default=None)\n    parser.add_argument(\"--transfer_mapping_only\", type=bool, default=False)\n    parser.add_argument(\"--start_iter\", type=int, default=0)\n    parser.add_argument(\"--iter\", type=int, default=20_000)\n\n    # model options\n    parser.add_argument(\"--size\", type=int, default=1024)\n    parser.add_argument(\"--min_rgb_size\", type=int, default=4)\n    parser.add_argument(\"--latent_size\", type=int, default=512)\n    parser.add_argument(\"--n_mlp\", type=int, default=8)\n    parser.add_argument(\"--n_sample\", type=int, default=60)\n    parser.add_argument(\"--constant_input\", type=bool, default=False)\n    parser.add_argument(\"--channel_multiplier\", type=int, default=2)\n    parser.add_argument(\"--d_skip\", type=bool, default=True)\n\n    # optimizer options\n    parser.add_argument(\"--lr\", type=float, default=0.002)\n    parser.add_argument(\"--d_lr_ratio\", type=float, default=1.0)\n    parser.add_argument(\"--lookahead\", type=bool, default=True)\n    parser.add_argument(\"--la_steps\", type=float, default=500)\n    parser.add_argument(\"--la_alpha\", type=float, default=0.5)\n\n    # loss options\n    parser.add_argument(\"--r1\", type=float, default=1e-5)\n    parser.add_argument(\"--path_regularize\", type=float, default=1)\n    parser.add_argument(\"--path_batch_shrink\", type=int, default=2)\n    parser.add_argument(\"--d_reg_every\", type=int, default=16)\n    parser.add_argument(\"--g_reg_every\", type=int, default=4)\n    parser.add_argument(\"--mixing_prob\", type=float, default=0.666)\n\n    # augmentation options\n    parser.add_argument(\"--augment\", type=bool, default=True)\n    parser.add_argument(\"--contrastive\", type=float, default=0)\n    parser.add_argument(\"--balanced_consistency\", type=float, default=0)\n    parser.add_argument(\"--augment_p\", type=float, default=0)\n    parser.add_argument(\"--ada_target\", type=float, default=0.6)\n    parser.add_argument(\"--ada_length\", type=int, default=15_000)\n\n    # validation options\n    parser.add_argument(\"--val_batch_size\", type=int, default=6)\n    parser.add_argument(\"--fid_n_sample\", type=int, default=2500)\n    parser.add_argument(\"--fid_truncation\", type=float, default=None)\n    parser.add_argument(\"--ppl_space\", choices=[\"z\", \"w\"], default=\"w\")\n    parser.add_argument(\"--ppl_n_sample\", type=int, default=1250)\n    parser.add_argument(\"--ppl_crop\", type=bool, default=False)\n\n    # logging options\n    parser.add_argument(\"--log_spec_norm\", type=bool, default=False)\n    parser.add_argument(\"--img_every\", type=int, default=1000)\n    parser.add_argument(\"--eval_every\", type=int, default=-1)\n    parser.add_argument(\"--checkpoint_every\", type=int, default=1000)\n    parser.add_argument(\"--profile_mem\", action=\"store_true\")\n\n    # (multi-)GPU options\n    parser.add_argument(\"--local_rank\", type=int, default=0)\n    parser.add_argument(\"--cudnn_benchmark\", type=bool, default=True)\n\n    args = parser.parse_args()\n    if args.balanced_consistency > 0 or args.contrastive > 0:\n        args.augment = True\n    args.name = os.path.splitext(os.path.basename(args.path))[0]\n    args.r1 = args.r1 * args.size ** 2\n\n    args.num_gpus = int(os.environ[\"WORLD_SIZE\"]) if \"WORLD_SIZE\" in os.environ else 1\n    th.backends.cudnn.benchmark = args.cudnn_benchmark\n    args.distributed = args.num_gpus > 1\n\n    # code for updating wandb configs that were incorrect\n    # if args.local_rank == 0:\n    #     api = wandb.Api()\n    #     run = api.run(\"wav/temperatuur/7kp6g0zt\")\n    #     run.config = vars(args)\n    #     run.update()\n    # exit()\n\n    if args.distributed:\n        th.cuda.set_device(args.local_rank)\n        th.distributed.init_process_group(backend=\"nccl\", init_method=\"env://\")\n        synchronize()\n\n    generator = Generator(\n        args.size,\n        args.latent_size,\n        args.n_mlp,\n        channel_multiplier=args.channel_multiplier,\n        constant_input=args.constant_input,\n        min_rgb_size=args.min_rgb_size,\n    ).to(device, non_blocking=True)\n    discriminator = Discriminator(args.size, channel_multiplier=args.channel_multiplier, use_skip=args.d_skip).to(\n        device\n    )\n\n    if args.log_spec_norm:\n        for name, parameter in generator.named_parameters():\n            if \"weight\" in name and parameter.squeeze().dim() > 1:\n                mod = generator\n                for attr in name.replace(\".weight\", \"\").split(\".\"):\n                    mod = getattr(mod, attr)\n                validation.track_spectral_norm(mod)\n        for name, parameter in discriminator.named_parameters():\n            if \"weight\" in name and parameter.squeeze().dim() > 1:\n                mod = discriminator\n                for attr in name.replace(\".weight\", \"\").split(\".\"):\n                    mod = getattr(mod, attr)\n                validation.track_spectral_norm(mod)\n\n    g_ema = Generator(\n        args.size,\n        args.latent_size,\n        args.n_mlp,\n        channel_multiplier=args.channel_multiplier,\n        constant_input=args.constant_input,\n        min_rgb_size=args.min_rgb_size,\n    ).to(device, non_blocking=True)\n    g_ema.requires_grad_(False)\n    g_ema.eval()\n    accumulate(g_ema, generator, 0)\n\n    if args.contrastive > 0:\n        contrast_learner = ContrastiveLearner(\n            discriminator,\n            args.size,\n            augment_fn=nn.Sequential(\n                nn.ReflectionPad2d(int((math.sqrt(2) - 1) * args.size / 4)),  # zoom out\n                augs.RandomHorizontalFlip(),\n                RandomApply(augs.RandomAffine(degrees=0, translate=(0.25, 0.25), shear=(15, 15)), p=0.1),\n                RandomApply(augs.RandomRotation(180), p=0.1),\n                augs.RandomResizedCrop(size=(args.size, args.size), scale=(1, 1), ratio=(1, 1)),\n                RandomApply(augs.RandomResizedCrop(size=(args.size, args.size), scale=(0.5, 0.9)), p=0.1),  # zoom in\n                RandomApply(augs.RandomErasing(), p=0.1),\n            ),\n            hidden_layer=(-1, 0),\n        )\n    else:\n        contrast_learner = None\n\n    g_reg_ratio = args.g_reg_every / (args.g_reg_every + 1)\n    d_reg_ratio = args.d_reg_every / (args.d_reg_every + 1)\n\n    g_optim = th.optim.Adam(\n        generator.parameters(), lr=args.lr * g_reg_ratio, betas=(0 ** g_reg_ratio, 0.99 ** g_reg_ratio),\n    )\n    d_optim = th.optim.Adam(\n        discriminator.parameters(),\n        lr=args.lr * d_reg_ratio * args.d_lr_ratio,\n        betas=(0 ** d_reg_ratio, 0.99 ** d_reg_ratio),\n    )\n\n    if args.lookahead:\n        g_optim = LookaheadMinimax(\n            g_optim, d_optim, la_steps=args.la_steps, la_alpha=args.la_alpha, accumulate=args.num_accumulate\n        )\n\n    if args.checkpoint is not None:\n        print(\"load model:\", args.checkpoint)\n\n        checkpoint = th.load(args.checkpoint, map_location=lambda storage, loc: storage)\n\n        try:\n            ckpt_name = os.path.basename(args.checkpoint)\n            args.start_iter = int(os.path.splitext(ckpt_name)[-1].replace(args.name, \"\"))\n        except ValueError:\n            pass\n\n        if args.transfer_mapping_only:\n            print(\"Using generator latent mapping network from checkpoint\")\n            mapping_state_dict = {}\n            for key, val in checkpoint[\"state_dict\"].items():\n                if \"generator.style\" in key:\n                    mapping_state_dict[key.replace(\"generator.\", \"\")] = val\n            generator.load_state_dict(mapping_state_dict, strict=False)\n        else:\n            generator.load_state_dict(checkpoint[\"g\"], strict=False)\n            g_ema.load_state_dict(checkpoint[\"g_ema\"], strict=False)\n\n            discriminator.load_state_dict(checkpoint[\"d\"], strict=False)\n\n            if args.lookahead:\n                g_optim.load_state_dict(checkpoint[\"g_optim\"], checkpoint[\"d_optim\"])\n            else:\n                g_optim.load_state_dict(checkpoint[\"g_optim\"])\n                d_optim.load_state_dict(checkpoint[\"d_optim\"])\n\n        del checkpoint\n        th.cuda.empty_cache()\n\n    if args.distributed:\n        generator = nn.parallel.DistributedDataParallel(\n            generator,\n            device_ids=[args.local_rank],\n            output_device=args.local_rank,\n            broadcast_buffers=False,\n            find_unused_parameters=True,\n        )\n\n        discriminator = nn.parallel.DistributedDataParallel(\n            discriminator,\n            device_ids=[args.local_rank],\n            output_device=args.local_rank,\n            broadcast_buffers=False,\n            find_unused_parameters=True,\n        )\n\n        if contrast_learner is not None:\n            contrast_learner = nn.parallel.DistributedDataParallel(\n                contrast_learner,\n                device_ids=[args.local_rank],\n                output_device=args.local_rank,\n                broadcast_buffers=False,\n                find_unused_parameters=True,\n            )\n\n    transform = transforms.Compose(\n        [\n            transforms.RandomVerticalFlip(p=0.5 if args.vflip else 0),\n            transforms.RandomHorizontalFlip(p=0.5 if args.hflip else 0),\n            transforms.ToTensor(),\n            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n        ]\n    )\n\n    dataset = MultiResolutionDataset(args.path, transform, args.size)\n    loader = data.DataLoader(\n        dataset,\n        batch_size=args.batch_size,\n        sampler=data_sampler(dataset, shuffle=True, distributed=args.distributed),\n        num_workers=8,\n        drop_last=True,\n        pin_memory=True,\n    )\n\n    if get_rank() == 0:\n        validation.get_dataset_inception_features(loader, args.name, args.size)\n        if args.wbgroup is None:\n            wandb.init(project=args.wbproj, name=args.wbname, config=vars(args))\n        else:\n            wandb.init(project=args.wbproj, group=args.wbgroup, name=args.wbname, config=vars(args))\n\n    if args.profile_mem:\n        os.environ[\"GPU_DEBUG\"] = str(args.local_rank)\n        from gpu_profile import gpu_profile\n\n        sys.settrace(gpu_profile)\n\n    train(args, loader, generator, discriminator, contrast_learner, g_optim, d_optim, g_ema)\n"
  },
  {
    "path": "train_profile.py",
    "content": "import argparse\nimport gc\nimport math\nimport os\nimport random\nimport sys\nimport time\n\nimport numpy as np\nimport torch as th\nimport torch.autograd.profiler as profiler\nimport wandb\nfrom contrastive_learner import ContrastiveLearner, RandomApply\nfrom kornia import augmentation as augs\nfrom scipy.ndimage import gaussian_filter\nfrom torch import nn\nfrom torch.nn import functional as F\nfrom torch.utils import data\nfrom torchvision import transforms, utils\nfrom tqdm import tqdm\n\nimport validation\nfrom augment import augment\nfrom dataset import MultiResolutionDataset\nfrom distributed import get_rank, reduce_loss_dict, reduce_sum, synchronize\nfrom lookahead_minimax import LookaheadMinimax\nfrom models.stylegan2 import Discriminator, Generator\n\nsys.path.insert(0, \"../lookahead_minimax\")\n\n\ndef data_sampler(dataset, shuffle, distributed):\n    if distributed:\n        return data.distributed.DistributedSampler(dataset, shuffle=shuffle)\n    if shuffle:\n        return data.RandomSampler(dataset)\n    else:\n        return data.SequentialSampler(dataset)\n\n\ndef requires_grad(model, flag=True):\n    for p in model.parameters():\n        p.requires_grad = flag\n\n\ndef accumulate(model1, model2, decay=0.5 ** (32.0 / 10_000)):\n    par1 = dict(model1.named_parameters())\n    par2 = dict(model2.named_parameters())\n    for name, param in model1.named_parameters():\n        param.data = decay * par1[name].data + (1 - decay) * par2[name].data\n\n\ndef sample_data(loader):\n    while True:\n        for batch in loader:\n            yield batch\n\n\ndef make_noise(batch_size, latent_dim, prob):\n    if prob > 0 and random.random() < prob:\n        return th.randn(2, batch_size, latent_dim, device=device).unbind(0)\n    else:\n        return [th.randn(batch_size, latent_dim, device=device)]\n\n\ndef d_logistic_loss(real_pred, fake_pred):\n    real_loss = F.softplus(-real_pred)\n    fake_loss = F.softplus(fake_pred)\n    return real_loss.mean() + fake_loss.mean()\n\n\ndef d_r1_penalty(real_img, real_pred, args):\n    (grad_real,) = th.autograd.grad(real_pred.sum(), real_img, create_graph=True)\n    r1_loss = grad_real.pow(2).view(grad_real.shape[0], -1).sum(1).mean()\n    r1_loss = r1_loss / 2.0 + 0 * real_pred[0]\n    return r1_loss\n\n\ndef g_non_saturating_loss(fake_pred):\n    return F.softplus(-fake_pred).mean()\n\n\ndef g_path_length_regularization(generator, mean_path_length, args):\n    path_batch_size = max(1, args.batch_size // args.path_batch_shrink)\n\n    noise = make_noise(path_batch_size, args.latent_size, args.mixing_prob)\n    fake_img, latents = generator(noise, return_latents=True)\n\n    img_noise = th.randn_like(fake_img) / math.sqrt(fake_img.shape[2] * fake_img.shape[3])\n    noisy_img_sum = (fake_img * img_noise).sum()\n\n    (grad,) = th.autograd.grad(noisy_img_sum, latents, create_graph=True)\n\n    path_lengths = th.sqrt(grad.pow(2).sum(2).mean(1))\n    path_mean = mean_path_length + 0.01 * (path_lengths.mean() - mean_path_length)\n    path_loss = (path_lengths - path_mean).pow(2).mean()\n    if not th.isnan(path_mean):\n        mean_path_length = path_mean.detach()\n\n    if args.path_batch_shrink:\n        path_loss += 0 * fake_img[0, 0, 0, 0]\n\n    return path_loss, mean_path_length\n\n\n# detach / item all the things\n# separate out into smaller functions so those local scopes get cleaned better\n# placeholder tensor pattern\n# torch no grad where possible\ndef train(args, loader, generator, discriminator, contrast_learner, g_optim, d_optim, g_ema):\n    if args.distributed:\n        g_module = generator.module\n        d_module = discriminator.module\n        if contrast_learner is not None:\n            cl_module = contrast_learner.module\n    else:\n        g_module = generator\n        d_module = discriminator\n        cl_module = contrast_learner\n\n    loader = sample_data(loader)\n    sample_z = th.randn(args.n_sample, args.latent_size, device=device)\n    mse = th.nn.MSELoss()\n    mean_path_length = 0\n    ada_augment = th.tensor([0.0, 0.0], device=device)\n    ada_aug_p = args.augment_p if args.augment_p > 0 else 0.0\n    ada_aug_step = args.ada_target / args.ada_length\n    r_t_stat = 0\n    fids = []\n\n    pbar = range(args.iter)\n    if get_rank() == 0:\n        pbar = tqdm(pbar, initial=args.start_iter, dynamic_ncols=True, smoothing=0)\n\n    for idx in pbar:\n        i = idx + args.start_iter\n        if i > args.iter:\n            print(\"Done!\")\n            break\n\n        if idx == args.nsys_iter:\n            print(\"Profiling begun at iteration {}\".format(idx))\n            th.cuda.cudart().cudaProfilerStart()\n        if idx >= args.nsys_iter:\n            th.cuda.nvtx.range_push(\"Iter {}\".format(idx))\n\n        tick_start = time.time()\n\n        loss_dict = {\n            \"Generator\": th.tensor(0, device=device).float(),\n            \"Discriminator\": th.tensor(0, device=device).float(),\n            \"Real Score\": th.tensor(0, device=device).float(),\n            \"Fake Score\": th.tensor(0, device=device).float(),\n            \"Contrastive\": th.tensor(0, device=device).float(),\n            \"Consistency\": th.tensor(0, device=device).float(),\n            \"R1 Penalty\": th.tensor(0, device=device).float(),\n            \"Path Length Regularization\": th.tensor(0, device=device).float(),\n            \"Augment\": th.tensor(0, device=device).float(),\n            \"Rt\": th.tensor(0, device=device).float(),\n        }\n\n        with profiler.record_function(\"D train\"):\n            if idx >= args.nsys_iter:\n                th.cuda.nvtx.range_push(\"D train\")\n            requires_grad(generator, False)\n            requires_grad(discriminator, True)\n\n            discriminator.zero_grad()\n            for _ in range(args.num_accumulate):\n                real_img_og = next(loader).to(device, non_blocking=True)\n                noise = make_noise(args.batch_size, args.latent_size, args.mixing_prob)\n                fake_img_og, _ = generator(noise)\n                if args.augment:\n                    fake_img, _ = augment(fake_img_og, ada_aug_p)\n                    real_img, _ = augment(real_img_og, ada_aug_p)\n                else:\n                    fake_img = fake_img_og\n                    real_img = real_img_og\n\n                fake_pred = discriminator(fake_img)\n                real_pred = discriminator(real_img)\n                logistic_loss = d_logistic_loss(real_pred, fake_pred)\n                loss_dict[\"Discriminator\"] += logistic_loss.detach()\n                loss_dict[\"Real Score\"] += real_pred.mean().detach()\n                loss_dict[\"Fake Score\"] += fake_pred.mean().detach()\n                d_loss = logistic_loss\n\n                if args.contrastive > 0:\n                    contrast_learner(fake_img_og, fake_img, accumulate=True)\n                    contrast_learner(real_img_og, real_img, accumulate=True)\n                    contrast_loss = cl_module.calculate_loss()\n                    loss_dict[\"Contrastive\"] += contrast_loss.detach()\n                    d_loss += args.contrastive * contrast_loss\n\n                if args.balanced_consistency > 0:\n                    consistency_loss = mse(real_pred, discriminator(real_img_og)) + mse(\n                        fake_pred, discriminator(fake_img_og)\n                    )\n                    loss_dict[\"Consistency\"] += consistency_loss.detach()\n                    d_loss += args.balanced_consistency * consistency_loss\n\n                d_loss /= args.num_accumulate\n                d_loss.backward()\n            d_optim.step()\n            if idx >= args.nsys_iter:\n                th.cuda.nvtx.range_pop()\n\n        if args.r1 > 0 and i % args.d_reg_every == 0:\n            with profiler.record_function(\"D reg\"):\n                if idx >= args.nsys_iter:\n                    th.cuda.nvtx.range_push(\"D reg\")\n                discriminator.zero_grad()\n                for _ in range(args.num_accumulate):\n                    real_img = next(loader).to(device, non_blocking=True)\n                    real_img.requires_grad = True\n                    real_pred = discriminator(real_img)\n                    r1_loss = d_r1_penalty(real_img, real_pred, args)\n                    loss_dict[\"R1 Penalty\"] += r1_loss.detach().squeeze()\n                    r1_loss = args.r1 * args.d_reg_every * r1_loss / args.num_accumulate\n                    r1_loss.backward()\n                d_optim.step()\n                if idx >= args.nsys_iter:\n                    th.cuda.nvtx.range_pop()\n\n        if args.augment and args.augment_p == 0:\n            with profiler.record_function(\"ADA\"):\n                if idx >= args.nsys_iter:\n                    th.cuda.nvtx.range_push(\"ADA\")\n                ada_augment += th.tensor((th.sign(real_pred).sum().item(), real_pred.shape[0]), device=device)\n                ada_augment = reduce_sum(ada_augment)\n\n                if ada_augment[1] > 255:\n                    pred_signs, n_pred = ada_augment.tolist()\n\n                    r_t_stat = pred_signs / n_pred\n                    loss_dict[\"Rt\"] = th.tensor(r_t_stat, device=device).float()\n                    if r_t_stat > args.ada_target:\n                        sign = 1\n                    else:\n                        sign = -1\n\n                    ada_aug_p += sign * ada_aug_step * n_pred\n                    ada_aug_p = min(1, max(0, ada_aug_p))\n                    ada_augment.mul_(0)\n                    loss_dict[\"Augment\"] = th.tensor(ada_aug_p, device=device).float()\n                if idx >= args.nsys_iter:\n                    th.cuda.nvtx.range_pop()\n\n        with profiler.record_function(\"G train\"):\n            if idx >= args.nsys_iter:\n                th.cuda.nvtx.range_push(\"G train\")\n            requires_grad(generator, True)\n            requires_grad(discriminator, False)\n\n            generator.zero_grad()\n            for _ in range(args.num_accumulate):\n                noise = make_noise(args.batch_size, args.latent_size, args.mixing_prob)\n                fake_img, _ = generator(noise)\n                if args.augment:\n                    fake_img, _ = augment(fake_img, ada_aug_p)\n                fake_pred = discriminator(fake_img)\n                g_loss = g_non_saturating_loss(fake_pred)\n                loss_dict[\"Generator\"] += g_loss.detach()\n                g_loss /= args.num_accumulate\n                g_loss.backward()\n            g_optim.step()\n            if idx >= args.nsys_iter:\n                th.cuda.nvtx.range_pop()\n\n        if args.path_regularize > 0 and i % args.g_reg_every == 0:\n            with profiler.record_function(\"G reg\"):\n                if idx >= args.nsys_iter:\n                    th.cuda.nvtx.range_push(\"G reg\")\n                generator.zero_grad()\n                for _ in range(args.num_accumulate):\n                    path_loss, mean_path_length = g_path_length_regularization(generator, mean_path_length, args)\n                    loss_dict[\"Path Length Regularization\"] += path_loss.detach()\n                    path_loss = args.path_regularize * args.g_reg_every * path_loss / args.num_accumulate\n                    path_loss.backward()\n                g_optim.step()\n                if idx >= args.nsys_iter:\n                    th.cuda.nvtx.range_pop()\n\n        with profiler.record_function(\"Log / Eval\"):\n            if idx >= args.nsys_iter:\n                th.cuda.nvtx.range_push(\"G accum\")\n            accumulate(g_ema, g_module)\n            if idx >= args.nsys_iter:\n                th.cuda.nvtx.range_pop()\n\n            if idx >= args.nsys_iter:\n                th.cuda.nvtx.range_push(\"Log / Eval\")\n            loss_reduced = reduce_loss_dict(loss_dict)\n            log_dict = {k: v.mean().item() / args.num_accumulate for k, v in loss_reduced.items() if v != 0}\n            log_dict[\"Tick Length\"] = time.time() - tick_start\n\n            if get_rank() == 0:\n                with th.no_grad():\n                    if args.log_spec_norm:\n                        G_norms = []\n                        for name, spec_norm in g_module.named_buffers():\n                            if \"spectral_norm\" in name:\n                                G_norms.append(spec_norm.cpu().numpy())\n                        G_norms = np.array(G_norms)\n                        D_norms = []\n                        for name, spec_norm in d_module.named_buffers():\n                            if \"spectral_norm\" in name:\n                                D_norms.append(spec_norm.cpu().numpy())\n                        D_norms = np.array(D_norms)\n                        log_dict[f\"Spectral Norms/G min spectral norm\"] = np.log(G_norms).min()\n                        log_dict[f\"Spectral Norms/G mean spectral norm\"] = np.log(G_norms).mean()\n                        log_dict[f\"Spectral Norms/G max spectral norm\"] = np.log(G_norms).max()\n                        log_dict[f\"Spectral Norms/D min spectral norm\"] = np.log(D_norms).min()\n                        log_dict[f\"Spectral Norms/D mean spectral norm\"] = np.log(D_norms).mean()\n                        log_dict[f\"Spectral Norms/D max spectral norm\"] = np.log(D_norms).max()\n\n                    if args.img_every != -1 and i % args.img_every == 0:\n                        g_ema.eval()\n                        sample = []\n                        for sub in range(0, len(sample_z), args.batch_size):\n                            subsample, _ = g_ema([sample_z[sub : sub + args.batch_size]])\n                            sample.append(subsample.detach().cpu())\n                        sample = th.cat(sample)\n                        grid = utils.make_grid(sample, nrow=10, normalize=True, range=(-1, 1))\n                        log_dict[\"Generated Images EMA\"] = [wandb.Image(grid, caption=f\"Step {i}\")]\n\n                    if args.eval_every != -1 and i % args.eval_every == 0:\n                        fid_dict = validation.fid(\n                            g_ema, args.val_batch_size, args.fid_n_sample, args.fid_truncation, args.name\n                        )\n\n                        fid = fid_dict[\"FID\"]\n                        fids.append(fid)\n                        density = fid_dict[\"Density\"]\n                        coverage = fid_dict[\"Coverage\"]\n\n                        ppl = validation.ppl(\n                            g_ema,\n                            args.val_batch_size,\n                            args.ppl_n_sample,\n                            args.ppl_space,\n                            args.ppl_crop,\n                            args.latent_size,\n                        )\n\n                        log_dict[\"Evaluation/FID\"] = fid\n                        log_dict[\"Sweep/FID_smooth\"] = gaussian_filter(np.array(fids), [5])[-1]\n                        log_dict[\"Evaluation/Density\"] = density\n                        log_dict[\"Evaluation/Coverage\"] = coverage\n                        log_dict[\"Evaluation/PPL\"] = ppl\n\n                        gc.collect()\n                        th.cuda.empty_cache()\n\n                    wandb.log(log_dict)\n\n                    if args.eval_every != -1:\n                        description = (\n                            f\"FID: {fid:.4f}   PPL: {ppl:.4f}   Dens: {density:.4f}   Cov: {coverage:.4f}   \"\n                            + f\"G: {log_dict['Generator']:.4f}   D: {log_dict['Discriminator']:.4f}\"\n                        )\n                    else:\n                        description = f\"G: {log_dict['Generator']:.4f}   D: {log_dict['Discriminator']:.4f}\"\n                    if \"Augment\" in log_dict:\n                        description += f\"   Aug: {log_dict['Augment']:.4f}\"  #   Rt: {log_dict['Rt']:.4f}\"\n                    if \"R1 Penalty\" in log_dict:\n                        description += f\"   R1: {log_dict['R1 Penalty']:.4f}\"\n                    if \"Path Length Regularization\" in log_dict:\n                        description += f\"   Path: {log_dict['Path Length Regularization']:.4f}\"\n                    pbar.set_description(description)\n\n                    if i % args.checkpoint_every == 0:\n                        check_name = \"-\".join(\n                            [\n                                args.name,\n                                args.wbname,\n                                wandb.run.dir.split(\"/\")[-1].split(\"-\")[-1],\n                                # str(int(fid)),\n                                str(args.size),\n                                str(i).zfill(6),\n                            ]\n                        )\n                        th.save(\n                            {\n                                \"g\": g_module.state_dict(),\n                                \"d\": d_module.state_dict(),\n                                # \"cl\": cl_module.state_dict(),\n                                \"g_ema\": g_ema.state_dict(),\n                                \"g_optim\": g_optim.state_dict(),\n                                \"d_optim\": d_optim.state_dict(),\n                            },\n                            f\"/home/hans/modelzoo/maua-sg2/{check_name}.pt\",\n                        )\n\n            if idx >= args.nsys_iter:\n                th.cuda.nvtx.range_pop()\n            if idx >= args.nsys_iter:\n                th.cuda.nvtx.range_pop()  # iteration range\n\n        gpu_profile(frame=sys._getframe(), event=\"line\", arg=None)\n\n    if args.nsys_iter != -1:\n        th.cuda.cudart().cudaProfilerStop()\n\n\nif __name__ == \"__main__\":\n    device = \"cuda\"\n\n    parser = argparse.ArgumentParser()\n\n    parser.add_argument(\"--wbname\", type=str, required=True)\n    parser.add_argument(\"--wbproj\", type=str, required=True)\n    parser.add_argument(\"--wbgroup\", type=str, default=None)\n\n    # data options\n    parser.add_argument(\"--path\", type=str, required=True)\n    parser.add_argument(\"--vflip\", type=bool, default=False)\n    parser.add_argument(\"--hflip\", type=bool, default=True)\n\n    # training options\n    parser.add_argument(\"--batch_size\", type=int, default=12)\n    parser.add_argument(\"--num_accumulate\", type=int, default=1)\n\n    parser.add_argument(\"--checkpoint\", type=str, default=None)\n    parser.add_argument(\"--transfer_mapping_only\", type=bool, default=False)\n    parser.add_argument(\"--start_iter\", type=int, default=0)\n    parser.add_argument(\"--iter\", type=int, default=60_000)\n\n    # model options\n    parser.add_argument(\"--size\", type=int, default=256)\n    parser.add_argument(\"--min_rgb_size\", type=int, default=4)\n    parser.add_argument(\"--latent_size\", type=int, default=512)\n    parser.add_argument(\"--n_mlp\", type=int, default=8)\n    parser.add_argument(\"--n_sample\", type=int, default=60)\n    parser.add_argument(\"--constant_input\", type=bool, default=False)\n    parser.add_argument(\"--channel_multiplier\", type=int, default=2)\n    parser.add_argument(\"--d_skip\", type=bool, default=True)\n\n    # optimizer options\n    parser.add_argument(\"--lr\", type=float, default=0.002)\n    parser.add_argument(\"--d_lr_ratio\", type=float, default=1.0)\n    parser.add_argument(\"--lookahead\", type=bool, default=True)\n    parser.add_argument(\"--la_steps\", type=float, default=500)\n    parser.add_argument(\"--la_alpha\", type=float, default=0.5)\n\n    # loss options\n    parser.add_argument(\"--r1\", type=float, default=1e-5)\n    parser.add_argument(\"--path_regularize\", type=float, default=1)\n    parser.add_argument(\"--path_batch_shrink\", type=int, default=2)\n    parser.add_argument(\"--d_reg_every\", type=int, default=16)\n    parser.add_argument(\"--g_reg_every\", type=int, default=4)\n    parser.add_argument(\"--mixing_prob\", type=float, default=0.666)\n\n    # augmentation options\n    parser.add_argument(\"--augment\", type=bool, default=False)\n    parser.add_argument(\"--contrastive\", type=float, default=0)\n    parser.add_argument(\"--balanced_consistency\", type=float, default=0)\n    parser.add_argument(\"--augment_p\", type=float, default=0)\n    parser.add_argument(\"--ada_target\", type=float, default=0.6)\n    parser.add_argument(\"--ada_length\", type=int, default=40_000)\n\n    # validation options\n    parser.add_argument(\"--val_batch_size\", type=int, default=6)\n    parser.add_argument(\"--fid_n_sample\", type=int, default=2500)\n    parser.add_argument(\"--fid_truncation\", type=float, default=None)\n    parser.add_argument(\"--ppl_space\", choices=[\"z\", \"w\"], default=\"w\")\n    parser.add_argument(\"--ppl_n_sample\", type=int, default=1250)\n    parser.add_argument(\"--ppl_crop\", type=bool, default=False)\n\n    # logging options\n    parser.add_argument(\"--log_spec_norm\", type=bool, default=False)\n    parser.add_argument(\"--img_every\", type=int, default=1000)\n    parser.add_argument(\"--eval_every\", type=int, default=1000)\n    parser.add_argument(\"--checkpoint_every\", type=int, default=1000)\n\n    # (multi-)GPU options\n    parser.add_argument(\"--local_rank\", type=int, default=0)\n    parser.add_argument(\"--cudnn_benchmark\", type=bool, default=True)\n    parser.add_argument(\"--nsys_iter\", type=int, default=-1)\n    parser.add_argument(\"--th_prof\", action=\"store_true\")\n    parser.add_argument(\"--prof_gpu\", action=\"store_true\")\n\n    args = parser.parse_args()\n\n    with th.autograd.profiler.profile(\n        enabled=args.th_prof, use_cuda=True, record_shapes=True, profile_memory=True, with_stack=True\n    ) as prof:\n        with profiler.record_function(\"init\"):\n            if args.balanced_consistency > 0 or args.contrastive > 0:\n                args.augment = True\n            args.name = os.path.splitext(os.path.basename(args.path))[0]\n            args.r1 = args.r1 * args.size ** 2\n\n            args.num_gpus = int(os.environ[\"WORLD_SIZE\"]) if \"WORLD_SIZE\" in os.environ else 1\n            th.backends.cudnn.benchmark = args.cudnn_benchmark\n            args.distributed = args.num_gpus > 1\n\n            if args.distributed:\n                th.cuda.set_device(args.local_rank)\n                th.distributed.init_process_group(backend=\"nccl\", init_method=\"env://\")\n                synchronize()\n\n            generator = Generator(\n                args.size,\n                args.latent_size,\n                args.n_mlp,\n                channel_multiplier=args.channel_multiplier,\n                constant_input=args.constant_input,\n                min_rgb_size=args.min_rgb_size,\n            ).to(device, non_blocking=True)\n            discriminator = Discriminator(\n                args.size, channel_multiplier=args.channel_multiplier, use_skip=args.d_skip\n            ).to(device, non_blocking=True)\n\n            if args.log_spec_norm:\n                for name, parameter in generator.named_parameters():\n                    if \"weight\" in name and parameter.squeeze().dim() > 1:\n                        mod = generator\n                        for attr in name.replace(\".weight\", \"\").split(\".\"):\n                            mod = getattr(mod, attr)\n                        validation.track_spectral_norm(mod)\n                for name, parameter in discriminator.named_parameters():\n                    if \"weight\" in name and parameter.squeeze().dim() > 1:\n                        mod = discriminator\n                        for attr in name.replace(\".weight\", \"\").split(\".\"):\n                            mod = getattr(mod, attr)\n                        validation.track_spectral_norm(mod)\n\n            g_ema = Generator(\n                args.size,\n                args.latent_size,\n                args.n_mlp,\n                channel_multiplier=args.channel_multiplier,\n                constant_input=args.constant_input,\n                min_rgb_size=args.min_rgb_size,\n            ).to(device, non_blocking=True)\n            g_ema.requires_grad_(False)\n            g_ema.eval()\n            accumulate(g_ema, generator, 0)\n\n            if args.contrastive > 0:\n                contrast_learner = ContrastiveLearner(\n                    discriminator,\n                    args.size,\n                    augment_fn=nn.Sequential(\n                        nn.ReflectionPad2d(int((math.sqrt(2) - 1) * args.size / 4)),  # zoom out\n                        augs.RandomHorizontalFlip(),\n                        RandomApply(augs.RandomAffine(degrees=0, translate=(0.25, 0.25), shear=(15, 15)), p=0.1),\n                        RandomApply(augs.RandomRotation(180), p=0.1),\n                        augs.RandomResizedCrop(size=(args.size, args.size), scale=(1, 1), ratio=(1, 1)),\n                        RandomApply(\n                            augs.RandomResizedCrop(size=(args.size, args.size), scale=(0.5, 0.9)), p=0.1\n                        ),  # zoom in\n                        RandomApply(augs.RandomErasing(), p=0.1),\n                    ),\n                    hidden_layer=(-1, 0),\n                )\n            else:\n                contrast_learner = None\n\n            g_reg_ratio = args.g_reg_every / (args.g_reg_every + 1)\n            d_reg_ratio = args.d_reg_every / (args.d_reg_every + 1)\n\n            g_optim = th.optim.Adam(\n                generator.parameters(), lr=args.lr * g_reg_ratio, betas=(0 ** g_reg_ratio, 0.99 ** g_reg_ratio),\n            )\n            d_optim = th.optim.Adam(\n                discriminator.parameters(),\n                lr=args.lr * d_reg_ratio * args.d_lr_ratio,\n                betas=(0 ** d_reg_ratio, 0.99 ** d_reg_ratio),\n            )\n\n            if args.lookahead:\n                g_optim = LookaheadMinimax(\n                    g_optim, d_optim, la_steps=args.la_steps, la_alpha=args.la_alpha, accumulate=args.num_accumulate\n                )\n\n            if args.checkpoint is not None:\n                print(\"load model:\", args.checkpoint)\n\n                checkpoint = th.load(args.checkpoint, map_location=lambda storage, loc: storage)\n\n                try:\n                    ckpt_name = os.path.basename(args.checkpoint)\n                    args.start_iter = int(os.path.splitext(ckpt_name)[-1].replace(args.name, \"\"))\n                except ValueError:\n                    pass\n\n                if args.transfer_mapping_only:\n                    print(\"Using generator latent mapping network from checkpoint\")\n                    mapping_state_dict = {}\n                    for key, val in checkpoint[\"state_dict\"].items():\n                        if \"generator.style\" in key:\n                            mapping_state_dict[key.replace(\"generator.\", \"\")] = val\n                    generator.load_state_dict(mapping_state_dict, strict=False)\n                else:\n                    generator.load_state_dict(checkpoint[\"g\"], strict=False)\n                    g_ema.load_state_dict(checkpoint[\"g_ema\"], strict=False)\n\n                    discriminator.load_state_dict(checkpoint[\"d\"], strict=False)\n\n                    if args.lookahead:\n                        g_optim.load_state_dict(checkpoint[\"g_optim\"], checkpoint[\"d_optim\"])\n                    else:\n                        g_optim.load_state_dict(checkpoint[\"g_optim\"])\n                        d_optim.load_state_dict(checkpoint[\"d_optim\"])\n\n                del checkpoint\n                th.cuda.empty_cache()\n\n            if args.distributed:\n                generator = nn.parallel.DistributedDataParallel(\n                    generator,\n                    device_ids=[args.local_rank],\n                    output_device=args.local_rank,\n                    broadcast_buffers=False,\n                    find_unused_parameters=True,\n                )\n\n                discriminator = nn.parallel.DistributedDataParallel(\n                    discriminator,\n                    device_ids=[args.local_rank],\n                    output_device=args.local_rank,\n                    broadcast_buffers=False,\n                    find_unused_parameters=True,\n                )\n\n                if contrast_learner is not None:\n                    contrast_learner = nn.parallel.DistributedDataParallel(\n                        contrast_learner,\n                        device_ids=[args.local_rank],\n                        output_device=args.local_rank,\n                        broadcast_buffers=False,\n                        find_unused_parameters=True,\n                    )\n\n            transform = transforms.Compose(\n                [\n                    transforms.RandomVerticalFlip(p=0.5 if args.vflip else 0),\n                    transforms.RandomHorizontalFlip(p=0.5 if args.hflip else 0),\n                    transforms.ToTensor(),\n                    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n                ]\n            )\n\n            dataset = MultiResolutionDataset(args.path, transform, args.size)\n            loader = data.DataLoader(\n                dataset,\n                batch_size=args.batch_size,\n                sampler=data_sampler(dataset, shuffle=True, distributed=args.distributed),\n                num_workers=0,\n                drop_last=True,\n                pin_memory=True,\n            )\n\n            if get_rank() == 0:\n                validation.get_dataset_inception_features(loader, args.name, args.size)\n                if args.wbgroup is None:\n                    wandb.init(project=args.wbproj, name=args.wbname, config=vars(args))\n                else:\n                    wandb.init(project=args.wbproj, group=args.wbgroup, name=args.wbname, config=vars(args))\n\n        if args.prof_gpu:\n            os.environ[\"GPU_DEBUG\"] = str(args.local_rank)\n            import sys\n            from gpu_profile import gpu_profile\n\n            sys.settrace(gpu_profile)\n        train(args, loader, generator, discriminator, contrast_learner, g_optim, d_optim, g_ema)\n\n    if args.th_prof:\n        print(prof.total_average())\n        print(\"cuda_memory_usage\", prof.table(sort_by=\"cuda_memory_usage\", row_limit=20))\n        prof.export_chrome_trace(f\"{args.name}_gpu{args.local_rank}.trace\")\n"
  },
  {
    "path": "validation/__init__.py",
    "content": "from .metrics import vae_fid, fid, get_dataset_inception_features, ppl, prdc\nfrom .spectral_norm import track_spectral_norm\n"
  },
  {
    "path": "validation/calc_fid.py",
    "content": "import argparse\nimport pickle\n\nimport torch\nfrom torch import nn\nimport numpy as np\nfrom scipy import linalg\nfrom tqdm import tqdm\n\nfrom model import Generator\nfrom inception import InceptionV3\n\n\n@torch.no_grad()\ndef extract_feature_from_samples(generator, inception, truncation, truncation_latent, batch_size, n_sample, device):\n    n_batch = n_sample // batch_size\n    resid = n_sample - (n_batch * batch_size)\n    batch_sizes = [batch_size] * n_batch + [resid]\n    features = []\n\n    for batch in tqdm(batch_sizes):\n        latent = torch.randn(batch, 512, device=device)\n        img, _ = g([latent], truncation=truncation, truncation_latent=truncation_latent)\n        feat = inception(img)[0].view(img.shape[0], -1)\n        features.append(feat.to(\"cpu\"))\n\n    features = torch.cat(features, 0)\n\n    return features\n\n\ndef calc_fid(sample_mean, sample_cov, real_mean, real_cov, eps=1e-6):\n    cov_sqrt, _ = linalg.sqrtm(sample_cov @ real_cov, disp=False)\n\n    if not np.isfinite(cov_sqrt).all():\n        print(\"product of cov matrices is singular\")\n        offset = np.eye(sample_cov.shape[0]) * eps\n        cov_sqrt = linalg.sqrtm((sample_cov + offset) @ (real_cov + offset))\n\n    if np.iscomplexobj(cov_sqrt):\n        if not np.allclose(np.diagonal(cov_sqrt).imag, 0, atol=1e-3):\n            m = np.max(np.abs(cov_sqrt.imag))\n\n            raise ValueError(f\"Imaginary component {m}\")\n\n        cov_sqrt = cov_sqrt.real\n\n    mean_diff = sample_mean - real_mean\n    mean_norm = mean_diff @ mean_diff\n\n    trace = np.trace(sample_cov) + np.trace(real_cov) - 2 * np.trace(cov_sqrt)\n\n    fid = mean_norm + trace\n\n    return fid\n\n\nif __name__ == \"__main__\":\n    device = \"cuda\"\n\n    parser = argparse.ArgumentParser()\n\n    parser.add_argument(\"--truncation\", type=float, default=1)\n    parser.add_argument(\"--truncation_mean\", type=int, default=4096 * 8)\n    parser.add_argument(\"--batch\", type=int, default=64)\n    parser.add_argument(\"--n_sample\", type=int, default=50000)\n    parser.add_argument(\"--size\", type=int, default=256)\n    parser.add_argument(\"--inception\", type=str, default=None, required=True)\n    parser.add_argument(\"ckpt\", metavar=\"CHECKPOINT\")\n\n    args = parser.parse_args()\n\n    ckpt = torch.load(args.ckpt)\n\n    g = Generator(args.size, 512, 8).to(device)\n    g.load_state_dict(ckpt[\"g_ema\"])\n    g = nn.DataParallel(g)\n    g.eval()\n\n    if args.truncation < 1:\n        with torch.no_grad():\n            mean_latent = g.mean_latent(args.truncation_mean)\n\n    else:\n        mean_latent = None\n\n    inception = InceptionV3([3], normalize_input=False, init_weights=False)\n    inception = nn.DataParallel(inception).eval().cuda()\n\n    features = extract_feature_from_samples(\n        g, inception, args.truncation, mean_latent, args.batch, args.n_sample, device\n    ).numpy()\n    print(f\"extracted {features.shape[0]} features\")\n\n    sample_mean = np.mean(features, 0)\n    sample_cov = np.cov(features, rowvar=False)\n\n    with open(args.inception, \"rb\") as f:\n        embeds = pickle.load(f)\n        real_mean = embeds[\"mean\"]\n        real_cov = embeds[\"cov\"]\n\n    fid = calc_fid(sample_mean, sample_cov, real_mean, real_cov)\n\n    print(\"fid:\", fid)\n"
  },
  {
    "path": "validation/calc_inception.py",
    "content": "import argparse\nimport pickle\nimport os\n\nimport torch\nfrom torch import nn\nfrom torch.nn import functional as F\nfrom torch.utils.data import DataLoader\nfrom torchvision import transforms\nfrom torchvision.models import Inception3\nimport numpy as np\nfrom tqdm import tqdm\n\nfrom inception import InceptionV3\nfrom dataset import MultiResolutionDataset\n\n\nclass Inception3Feature(Inception3):\n    def forward(self, x):\n        if x.shape[2] != 299 or x.shape[3] != 299:\n            x = F.interpolate(x, size=(299, 299), mode=\"bilinear\", align_corners=True)\n\n        x = self.Conv2d_1a_3x3(x)  # 299 x 299 x 3\n        x = self.Conv2d_2a_3x3(x)  # 149 x 149 x 32\n        x = self.Conv2d_2b_3x3(x)  # 147 x 147 x 32\n        x = F.max_pool2d(x, kernel_size=3, stride=2)  # 147 x 147 x 64\n\n        x = self.Conv2d_3b_1x1(x)  # 73 x 73 x 64\n        x = self.Conv2d_4a_3x3(x)  # 73 x 73 x 80\n        x = F.max_pool2d(x, kernel_size=3, stride=2)  # 71 x 71 x 192\n\n        x = self.Mixed_5b(x)  # 35 x 35 x 192\n        x = self.Mixed_5c(x)  # 35 x 35 x 256\n        x = self.Mixed_5d(x)  # 35 x 35 x 288\n\n        x = self.Mixed_6a(x)  # 35 x 35 x 288\n        x = self.Mixed_6b(x)  # 17 x 17 x 768\n        x = self.Mixed_6c(x)  # 17 x 17 x 768\n        x = self.Mixed_6d(x)  # 17 x 17 x 768\n        x = self.Mixed_6e(x)  # 17 x 17 x 768\n\n        x = self.Mixed_7a(x)  # 17 x 17 x 768\n        x = self.Mixed_7b(x)  # 8 x 8 x 1280\n        x = self.Mixed_7c(x)  # 8 x 8 x 2048\n\n        x = F.avg_pool2d(x, kernel_size=8)  # 8 x 8 x 2048\n\n        return x.view(x.shape[0], x.shape[1])  # 1 x 1 x 2048\n\n\ndef load_patched_inception_v3():\n    # inception = inception_v3(pretrained=True)\n    # inception_feat = Inception3Feature()\n    # inception_feat.load_state_dict(inception.state_dict())\n    inception_feat = InceptionV3([3], normalize_input=False, init_weights=False)\n\n    return inception_feat\n\n\n@torch.no_grad()\ndef extract_features(loader, inception, device):\n    pbar = tqdm(loader)\n\n    feature_list = []\n\n    for img in pbar:\n        img = img.to(device)\n        feature = inception(img)[0].view(img.shape[0], -1)\n        feature_list.append(feature.to(\"cpu\"))\n\n    features = torch.cat(feature_list, 0)\n\n    return features\n\n\nif __name__ == \"__main__\":\n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n    parser = argparse.ArgumentParser(description=\"Calculate Inception v3 features for datasets\")\n    parser.add_argument(\"--size\", type=int, default=256)\n    parser.add_argument(\"--batch\", default=64, type=int, help=\"batch size\")\n    parser.add_argument(\"--n_sample\", type=int, default=50000)\n    parser.add_argument(\"--vflip\", action=\"store_true\")\n    parser.add_argument(\"--hflip\", action=\"store_true\")\n    parser.add_argument(\"path\", metavar=\"PATH\", help=\"path to datset lmdb file\")\n\n    args = parser.parse_args()\n\n    inception = load_patched_inception_v3()\n    inception = nn.DataParallel(inception).eval().to(device)\n\n    transform = transforms.Compose(\n        [\n            transforms.RandomVerticalFlip(p=0.5 if args.vflip else 0),\n            transforms.RandomHorizontalFlip(p=0.5 if args.hflip else 0),\n            transforms.ToTensor(),\n            transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),\n        ]\n    )\n\n    dset = MultiResolutionDataset(args.path, transform=transform, resolution=args.size)\n    loader = DataLoader(dset, batch_size=args.batch, num_workers=4)\n\n    features = extract_features(loader, inception, device).numpy()\n\n    features = features[: args.n_sample]\n\n    print(f\"extracted {features.shape[0]} features\")\n\n    mean = np.mean(features, 0)\n    cov = np.cov(features, rowvar=False)\n\n    name = os.path.splitext(os.path.basename(args.path))[0]\n\n    with open(f\"inception_{name}.pkl\", \"wb\") as f:\n        pickle.dump({\"mean\": mean, \"cov\": cov, \"size\": args.size, \"path\": args.path}, f)\n"
  },
  {
    "path": "validation/calc_ppl.py",
    "content": "import argparse\r\n\r\nimport torch\r\nfrom torch.nn import functional as F\r\nimport numpy as np\r\nfrom tqdm import tqdm\r\n\r\nimport lpips\r\nfrom model import Generator\r\n\r\n\r\ndef normalize(x):\r\n    return x / torch.sqrt(x.pow(2).sum(-1, keepdim=True))\r\n\r\n\r\ndef slerp(a, b, t):\r\n    a = normalize(a)\r\n    b = normalize(b)\r\n    d = (a * b).sum(-1, keepdim=True)\r\n    p = t * torch.acos(d)\r\n    c = normalize(b - d * a)\r\n    d = a * torch.cos(p) + c * torch.sin(p)\r\n\r\n    return normalize(d)\r\n\r\n\r\ndef lerp(a, b, t):\r\n    return a + (b - a) * t\r\n\r\n\r\nif __name__ == \"__main__\":\r\n    device = \"cuda\"\r\n\r\n    parser = argparse.ArgumentParser()\r\n\r\n    parser.add_argument(\"--space\", choices=[\"z\", \"w\"])\r\n    parser.add_argument(\"--batch\", type=int, default=64)\r\n    parser.add_argument(\"--n_sample\", type=int, default=5000)\r\n    parser.add_argument(\"--size\", type=int, default=256)\r\n    parser.add_argument(\"--eps\", type=float, default=1e-4)\r\n    parser.add_argument(\"--crop\", action=\"store_true\")\r\n    parser.add_argument(\"ckpt\", metavar=\"CHECKPOINT\")\r\n\r\n    args = parser.parse_args()\r\n\r\n    latent_dim = 512\r\n\r\n    ckpt = torch.load(args.ckpt)\r\n\r\n    g = Generator(args.size, latent_dim, 8).to(device)\r\n    g.load_state_dict(ckpt[\"g_ema\"])\r\n    g.eval()\r\n\r\n    percept = lpips.PerceptualLoss(model=\"net-lin\", net=\"vgg\", use_gpu=device.startswith(\"cuda\"))\r\n\r\n    distances = []\r\n\r\n    n_batch = args.n_sample // args.batch\r\n    resid = args.n_sample - (n_batch * args.batch)\r\n    batch_sizes = [args.batch] * n_batch + [resid]\r\n\r\n    with torch.no_grad():\r\n        for batch in tqdm(batch_sizes):\r\n            noise = g.make_noise()\r\n\r\n            inputs = torch.randn([batch * 2, latent_dim], device=device)\r\n            lerp_t = torch.rand(batch, device=device)\r\n\r\n            if args.space == \"w\":\r\n                latent = g.get_latent(inputs)\r\n                latent_t0, latent_t1 = latent[::2], latent[1::2]\r\n                latent_e0 = lerp(latent_t0, latent_t1, lerp_t[:, None])\r\n                latent_e1 = lerp(latent_t0, latent_t1, lerp_t[:, None] + args.eps)\r\n                latent_e = torch.stack([latent_e0, latent_e1], 1).view(*latent.shape)\r\n\r\n            image, _ = g([latent_e], input_is_latent=True, noise=noise)\r\n\r\n            if args.crop:\r\n                c = image.shape[2] // 8\r\n                image = image[:, :, c * 3 : c * 7, c * 2 : c * 6]\r\n\r\n            factor = image.shape[2] // 256\r\n\r\n            if factor > 1:\r\n                image = F.interpolate(image, size=(256, 256), mode=\"bilinear\", align_corners=False)\r\n\r\n            dist = percept(image[::2], image[1::2]).view(image.shape[0] // 2) / (args.eps ** 2)\r\n            distances.append(dist.to(\"cpu\").numpy())\r\n\r\n    distances = np.concatenate(distances, 0)\r\n\r\n    lo = np.percentile(distances, 1, interpolation=\"lower\")\r\n    hi = np.percentile(distances, 99, interpolation=\"higher\")\r\n    filtered_dist = np.extract(np.logical_and(lo <= distances, distances <= hi), distances)\r\n\r\n    print(\"ppl:\", filtered_dist.mean())\r\n"
  },
  {
    "path": "validation/inception.py",
    "content": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torchvision import models\n\ntry:\n    from torchvision.models.utils import load_state_dict_from_url\nexcept ImportError:\n    from torch.utils.model_zoo import load_url as load_state_dict_from_url\n\n# Inception weights ported to Pytorch from\n# http://download.tensorflow.org/models/image/imagenet/inception-2015-12-05.tgz\nFID_WEIGHTS_URL = (\n    \"https://github.com/mseitzer/pytorch-fid/releases/download/fid_weights/pt_inception-2015-12-05-6726825d.pth\"\n)\n\n\nclass InceptionV3(nn.Module):\n    \"\"\"Pretrained InceptionV3 network returning feature maps\"\"\"\n\n    # Index of default block of inception to return,\n    # corresponds to output of final average pooling\n    DEFAULT_BLOCK_INDEX = 3\n\n    # Maps feature dimensionality to their output blocks indices\n    BLOCK_INDEX_BY_DIM = {\n        64: 0,  # First max pooling features\n        192: 1,  # Second max pooling featurs\n        768: 2,  # Pre-aux classifier features\n        2048: 3,  # Final average pooling features\n    }\n\n    def __init__(\n        self,\n        output_blocks=[DEFAULT_BLOCK_INDEX],\n        resize_input=True,\n        normalize_input=True,\n        requires_grad=False,\n        init_weights=False,\n        use_fid_inception=False,\n    ):\n        \"\"\"Build pretrained InceptionV3\n\n        Parameters\n        ----------\n        output_blocks : list of int\n            Indices of blocks to return features of. Possible values are:\n                - 0: corresponds to output of first max pooling\n                - 1: corresponds to output of second max pooling\n                - 2: corresponds to output which is fed to aux classifier\n                - 3: corresponds to output of final average pooling\n        resize_input : bool\n            If true, bilinearly resizes input to width and height 299 before\n            feeding input to model. As the network without fully connected\n            layers is fully convolutional, it should be able to handle inputs\n            of arbitrary size, so resizing might not be strictly needed\n        normalize_input : bool\n            If true, scales the input from range (0, 1) to the range the\n            pretrained Inception network expects, namely (-1, 1)\n        requires_grad : bool\n            If true, parameters of the model require gradients. Possibly useful\n            for finetuning the network\n        use_fid_inception : bool\n            If true, uses the pretrained Inception model used in Tensorflow's\n            FID implementation. If false, uses the pretrained Inception model\n            available in torchvision. The FID Inception model has different\n            weights and a slightly different structure from torchvision's\n            Inception model. If you want to compute FID scores, you are\n            strongly advised to set this parameter to true to get comparable\n            results.\n        \"\"\"\n        super(InceptionV3, self).__init__()\n\n        self.resize_input = resize_input\n        self.normalize_input = normalize_input\n        self.output_blocks = sorted(output_blocks)\n        self.last_needed_block = max(output_blocks)\n\n        assert self.last_needed_block <= 3, \"Last possible output block index is 3\"\n\n        self.blocks = nn.ModuleList()\n\n        if use_fid_inception:\n            inception = fid_inception_v3()\n        else:\n            inception = models.inception_v3(pretrained=True)  # , init_weights=False)\n\n        # Block 0: input to maxpool1\n        block0 = [\n            inception.Conv2d_1a_3x3,\n            inception.Conv2d_2a_3x3,\n            inception.Conv2d_2b_3x3,\n            nn.MaxPool2d(kernel_size=3, stride=2),\n        ]\n        self.blocks.append(nn.Sequential(*block0))\n\n        # Block 1: maxpool1 to maxpool2\n        if self.last_needed_block >= 1:\n            block1 = [inception.Conv2d_3b_1x1, inception.Conv2d_4a_3x3, nn.MaxPool2d(kernel_size=3, stride=2)]\n            self.blocks.append(nn.Sequential(*block1))\n\n        # Block 2: maxpool2 to aux classifier\n        if self.last_needed_block >= 2:\n            block2 = [\n                inception.Mixed_5b,\n                inception.Mixed_5c,\n                inception.Mixed_5d,\n                inception.Mixed_6a,\n                inception.Mixed_6b,\n                inception.Mixed_6c,\n                inception.Mixed_6d,\n                inception.Mixed_6e,\n            ]\n            self.blocks.append(nn.Sequential(*block2))\n\n        # Block 3: aux classifier to final avgpool\n        if self.last_needed_block >= 3:\n            block3 = [\n                inception.Mixed_7a,\n                inception.Mixed_7b,\n                inception.Mixed_7c,\n                nn.AdaptiveAvgPool2d(output_size=(1, 1)),\n            ]\n            self.blocks.append(nn.Sequential(*block3))\n\n        for param in self.parameters():\n            param.requires_grad = requires_grad\n\n    def forward(self, inp):\n        \"\"\"Get Inception feature maps\n\n        Parameters\n        ----------\n        inp : torch.autograd.Variable\n            Input tensor of shape Bx3xHxW. Values are expected to be in\n            range (0, 1)\n\n        Returns\n        -------\n        List of torch.autograd.Variable, corresponding to the selected output\n        block, sorted ascending by index\n        \"\"\"\n        outp = []\n        x = inp\n\n        if self.resize_input:\n            x = F.interpolate(x, size=(299, 299), mode=\"bilinear\", align_corners=False)\n\n        if self.normalize_input:\n            x = 2 * x - 1  # Scale from range (0, 1) to range (-1, 1)\n\n        for idx, block in enumerate(self.blocks):\n            x = block(x)\n            if idx in self.output_blocks:\n                outp.append(x)\n\n            if idx == self.last_needed_block:\n                break\n\n        return outp\n\n\ndef fid_inception_v3():\n    \"\"\"Build pretrained Inception model for FID computation\n\n    The Inception model for FID computation uses a different set of weights\n    and has a slightly different structure than torchvision's Inception.\n\n    This method first constructs torchvision's Inception and then patches the\n    necessary parts that are different in the FID Inception model.\n    \"\"\"\n    inception = models.inception_v3(num_classes=1008, aux_logits=False, pretrained=False)\n    inception.Mixed_5b = FIDInceptionA(192, pool_features=32)\n    inception.Mixed_5c = FIDInceptionA(256, pool_features=64)\n    inception.Mixed_5d = FIDInceptionA(288, pool_features=64)\n    inception.Mixed_6b = FIDInceptionC(768, channels_7x7=128)\n    inception.Mixed_6c = FIDInceptionC(768, channels_7x7=160)\n    inception.Mixed_6d = FIDInceptionC(768, channels_7x7=160)\n    inception.Mixed_6e = FIDInceptionC(768, channels_7x7=192)\n    inception.Mixed_7b = FIDInceptionE_1(1280)\n    inception.Mixed_7c = FIDInceptionE_2(2048)\n\n    state_dict = load_state_dict_from_url(FID_WEIGHTS_URL, progress=True)\n    inception.load_state_dict(state_dict)\n    return inception\n\n\nclass FIDInceptionA(models.inception.InceptionA):\n    \"\"\"InceptionA block patched for FID computation\"\"\"\n\n    def __init__(self, in_channels, pool_features):\n        super(FIDInceptionA, self).__init__(in_channels, pool_features)\n\n    def forward(self, x):\n        branch1x1 = self.branch1x1(x)\n\n        branch5x5 = self.branch5x5_1(x)\n        branch5x5 = self.branch5x5_2(branch5x5)\n\n        branch3x3dbl = self.branch3x3dbl_1(x)\n        branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)\n        branch3x3dbl = self.branch3x3dbl_3(branch3x3dbl)\n\n        # Patch: Tensorflow's average pool does not use the padded zero's in\n        # its average calculation\n        branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1, count_include_pad=False)\n        branch_pool = self.branch_pool(branch_pool)\n\n        outputs = [branch1x1, branch5x5, branch3x3dbl, branch_pool]\n        return torch.cat(outputs, 1)\n\n\nclass FIDInceptionC(models.inception.InceptionC):\n    \"\"\"InceptionC block patched for FID computation\"\"\"\n\n    def __init__(self, in_channels, channels_7x7):\n        super(FIDInceptionC, self).__init__(in_channels, channels_7x7)\n\n    def forward(self, x):\n        branch1x1 = self.branch1x1(x)\n\n        branch7x7 = self.branch7x7_1(x)\n        branch7x7 = self.branch7x7_2(branch7x7)\n        branch7x7 = self.branch7x7_3(branch7x7)\n\n        branch7x7dbl = self.branch7x7dbl_1(x)\n        branch7x7dbl = self.branch7x7dbl_2(branch7x7dbl)\n        branch7x7dbl = self.branch7x7dbl_3(branch7x7dbl)\n        branch7x7dbl = self.branch7x7dbl_4(branch7x7dbl)\n        branch7x7dbl = self.branch7x7dbl_5(branch7x7dbl)\n\n        # Patch: Tensorflow's average pool does not use the padded zero's in\n        # its average calculation\n        branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1, count_include_pad=False)\n        branch_pool = self.branch_pool(branch_pool)\n\n        outputs = [branch1x1, branch7x7, branch7x7dbl, branch_pool]\n        return torch.cat(outputs, 1)\n\n\nclass FIDInceptionE_1(models.inception.InceptionE):\n    \"\"\"First InceptionE block patched for FID computation\"\"\"\n\n    def __init__(self, in_channels):\n        super(FIDInceptionE_1, self).__init__(in_channels)\n\n    def forward(self, x):\n        branch1x1 = self.branch1x1(x)\n\n        branch3x3 = self.branch3x3_1(x)\n        branch3x3 = [\n            self.branch3x3_2a(branch3x3),\n            self.branch3x3_2b(branch3x3),\n        ]\n        branch3x3 = torch.cat(branch3x3, 1)\n\n        branch3x3dbl = self.branch3x3dbl_1(x)\n        branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)\n        branch3x3dbl = [\n            self.branch3x3dbl_3a(branch3x3dbl),\n            self.branch3x3dbl_3b(branch3x3dbl),\n        ]\n        branch3x3dbl = torch.cat(branch3x3dbl, 1)\n\n        # Patch: Tensorflow's average pool does not use the padded zero's in\n        # its average calculation\n        branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1, count_include_pad=False)\n        branch_pool = self.branch_pool(branch_pool)\n\n        outputs = [branch1x1, branch3x3, branch3x3dbl, branch_pool]\n        return torch.cat(outputs, 1)\n\n\nclass FIDInceptionE_2(models.inception.InceptionE):\n    \"\"\"Second InceptionE block patched for FID computation\"\"\"\n\n    def __init__(self, in_channels):\n        super(FIDInceptionE_2, self).__init__(in_channels)\n\n    def forward(self, x):\n        branch1x1 = self.branch1x1(x)\n\n        branch3x3 = self.branch3x3_1(x)\n        branch3x3 = [\n            self.branch3x3_2a(branch3x3),\n            self.branch3x3_2b(branch3x3),\n        ]\n        branch3x3 = torch.cat(branch3x3, 1)\n\n        branch3x3dbl = self.branch3x3dbl_1(x)\n        branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)\n        branch3x3dbl = [\n            self.branch3x3dbl_3a(branch3x3dbl),\n            self.branch3x3dbl_3b(branch3x3dbl),\n        ]\n        branch3x3dbl = torch.cat(branch3x3dbl, 1)\n\n        # Patch: The FID Inception model uses max pooling instead of average\n        # pooling. This is likely an error in this specific Inception\n        # implementation, as other Inception models use average pooling here\n        # (which matches the description in the paper).\n        branch_pool = F.max_pool2d(x, kernel_size=3, stride=1, padding=1)\n        branch_pool = self.branch_pool(branch_pool)\n\n        outputs = [branch1x1, branch3x3, branch3x3dbl, branch_pool]\n        return torch.cat(outputs, 1)\n"
  },
  {
    "path": "validation/lpips/__init__.py",
    "content": "from __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport numpy as np\nfrom skimage.measure import compare_ssim\nimport torch\nfrom torch.autograd import Variable\n\nfrom . import dist_model\n\n\nclass PerceptualLoss(torch.nn.Module):\n    def __init__(\n        self, model=\"net-lin\", net=\"alex\", colorspace=\"rgb\", spatial=False, use_gpu=True, gpu_ids=[0]\n    ):  # VGG using our perceptually-learned weights (LPIPS metric)\n        # def __init__(self, model='net', net='vgg', use_gpu=True): # \"default\" way of using VGG as a perceptual loss\n        super(PerceptualLoss, self).__init__()\n        # print('Setting up Perceptual loss...')\n        self.use_gpu = use_gpu\n        self.spatial = spatial\n        self.gpu_ids = gpu_ids\n        self.model = dist_model.DistModel()\n        self.model.initialize(\n            model=model, net=net, use_gpu=use_gpu, colorspace=colorspace, spatial=self.spatial, gpu_ids=gpu_ids\n        )\n        # print('...[%s] initialized'%self.model.name())\n        # print('...Done')\n\n    def forward(self, pred, target, normalize=False):\n        \"\"\"\n        Pred and target are Variables.\n        If normalize is True, assumes the images are between [0,1] and then scales them between [-1,+1]\n        If normalize is False, assumes the images are already between [-1,+1]\n\n        Inputs pred and target are Nx3xHxW\n        Output pytorch Variable N long\n        \"\"\"\n\n        if normalize:\n            target = 2 * target - 1\n            pred = 2 * pred - 1\n\n        return self.model.forward(target, pred)\n"
  },
  {
    "path": "validation/lpips/base_model.py",
    "content": "import os\nimport torch\n\n\nclass BaseModel:\n    def __init__(self):\n        pass\n\n    def name(self):\n        return \"BaseModel\"\n\n    def initialize(self, use_gpu=True, gpu_ids=[0]):\n        self.use_gpu = use_gpu\n        self.gpu_ids = gpu_ids\n\n    def forward(self):\n        pass\n\n    def get_image_paths(self):\n        pass\n\n    def optimize_parameters(self):\n        pass\n\n    def get_current_visuals(self):\n        return self.input\n\n    def get_current_errors(self):\n        return {}\n\n    def save(self, label):\n        pass\n\n    # helper saving function that can be used by subclasses\n    def save_network(self, network, path, network_label, epoch_label):\n        save_filename = \"%s_net_%s.pth\" % (epoch_label, network_label)\n        save_path = os.path.join(path, save_filename)\n        torch.save(network.state_dict(), save_path)\n\n    # helper loading function that can be used by subclasses\n    def load_network(self, network, network_label, epoch_label):\n        save_filename = \"%s_net_%s.pth\" % (epoch_label, network_label)\n        save_path = os.path.join(self.save_dir, save_filename)\n        print(\"Loading network from %s\" % save_path)\n        network.load_state_dict(torch.load(save_path))\n\n    def update_learning_rate():\n        pass\n\n    def get_image_paths(self):\n        return self.image_paths\n\n    def save_done(self, flag=False):\n        np.save(os.path.join(self.save_dir, \"done_flag\"), flag)\n        np.savetxt(os.path.join(self.save_dir, \"done_flag\"), [flag,], fmt=\"%i\")\n\n"
  },
  {
    "path": "validation/lpips/dist_model.py",
    "content": "import numpy as np\nimport torch\nimport os\nfrom collections import OrderedDict\nfrom torch.autograd import Variable\nfrom .base_model import BaseModel\nfrom scipy.ndimage import zoom\nfrom tqdm import tqdm\n\n\nfrom . import networks_basic as networks\nfrom . import util\n\n\nclass DistModel(BaseModel):\n    def name(self):\n        return self.model_name\n\n    def initialize(\n        self,\n        model=\"net-lin\",\n        net=\"alex\",\n        colorspace=\"Lab\",\n        pnet_rand=False,\n        pnet_tune=False,\n        model_path=None,\n        use_gpu=True,\n        printNet=False,\n        spatial=False,\n        is_train=False,\n        lr=0.0001,\n        beta1=0.5,\n        version=\"0.1\",\n        gpu_ids=[0],\n    ):\n        \"\"\"\n        INPUTS\n            model - ['net-lin'] for linearly calibrated network\n                    ['net'] for off-the-shelf network\n                    ['L2'] for L2 distance in Lab colorspace\n                    ['SSIM'] for ssim in RGB colorspace\n            net - ['squeeze','alex','vgg']\n            model_path - if None, will look in weights/[NET_NAME].pth\n            colorspace - ['Lab','RGB'] colorspace to use for L2 and SSIM\n            use_gpu - bool - whether or not to use a GPU\n            printNet - bool - whether or not to print network architecture out\n            spatial - bool - whether to output an array containing varying distances across spatial dimensions\n            spatial_shape - if given, output spatial shape. if None then spatial shape is determined automatically via spatial_factor (see below).\n            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.\n            spatial_order - spline order of filter for upsampling in spatial mode, by default 1 (bilinear).\n            is_train - bool - [True] for training mode\n            lr - float - initial learning rate\n            beta1 - float - initial momentum term for adam\n            version - 0.1 for latest, 0.0 was original (with a bug)\n            gpu_ids - int array - [0] by default, gpus to use\n        \"\"\"\n        BaseModel.initialize(self, use_gpu=use_gpu, gpu_ids=gpu_ids)\n\n        self.model = model\n        self.net = net\n        self.is_train = is_train\n        self.spatial = spatial\n        self.gpu_ids = gpu_ids\n        self.model_name = \"%s [%s]\" % (model, net)\n\n        if self.model == \"net-lin\":  # pretrained net + linear layer\n            self.net = networks.PNetLin(\n                pnet_rand=pnet_rand,\n                pnet_tune=pnet_tune,\n                pnet_type=net,\n                use_dropout=True,\n                spatial=spatial,\n                version=version,\n                lpips=True,\n            )\n            kw = {}\n            if not use_gpu:\n                kw[\"map_location\"] = \"cpu\"\n            if model_path is None:\n                import inspect\n\n                model_path = os.path.abspath(\n                    os.path.join(inspect.getfile(self.initialize), \"..\", \"weights/v%s/%s.pth\" % (version, net))\n                )\n\n            if not is_train:\n                # print(\"Loading model from: %s\" % model_path)\n                self.net.load_state_dict(torch.load(model_path, **kw), strict=False)\n\n        elif self.model == \"net\":  # pretrained network\n            self.net = networks.PNetLin(pnet_rand=pnet_rand, pnet_type=net, lpips=False)\n        elif self.model in [\"L2\", \"l2\"]:\n            self.net = networks.L2(use_gpu=use_gpu, colorspace=colorspace)  # not really a network, only for testing\n            self.model_name = \"L2\"\n        elif self.model in [\"DSSIM\", \"dssim\", \"SSIM\", \"ssim\"]:\n            self.net = networks.DSSIM(use_gpu=use_gpu, colorspace=colorspace)\n            self.model_name = \"SSIM\"\n        else:\n            raise ValueError(\"Model [%s] not recognized.\" % self.model)\n\n        self.parameters = list(self.net.parameters())\n\n        if self.is_train:  # training mode\n            # extra network on top to go from distances (d0,d1) => predicted human judgment (h*)\n            self.rankLoss = networks.BCERankingLoss()\n            self.parameters += list(self.rankLoss.net.parameters())\n            self.lr = lr\n            self.old_lr = lr\n            self.optimizer_net = torch.optim.Adam(self.parameters, lr=lr, betas=(beta1, 0.999))\n        else:  # test mode\n            self.net.eval()\n\n        if use_gpu:\n            self.net.to(gpu_ids[0])\n            self.net = torch.nn.DataParallel(self.net, device_ids=gpu_ids)\n            if self.is_train:\n                self.rankLoss = self.rankLoss.to(device=gpu_ids[0])  # just put this on GPU0\n\n        if printNet:\n            print(\"---------- Networks initialized -------------\")\n            networks.print_network(self.net)\n            print(\"-----------------------------------------------\")\n\n    def forward(self, in0, in1, retPerLayer=False):\n        \"\"\" Function computes the distance between image patches in0 and in1\n        INPUTS\n            in0, in1 - torch.Tensor object of shape Nx3xXxY - image patch scaled to [-1,1]\n        OUTPUT\n            computed distances between in0 and in1\n        \"\"\"\n\n        return self.net.forward(in0, in1, retPerLayer=retPerLayer)\n\n    # ***** TRAINING FUNCTIONS *****\n    def optimize_parameters(self):\n        self.forward_train()\n        self.optimizer_net.zero_grad()\n        self.backward_train()\n        self.optimizer_net.step()\n        self.clamp_weights()\n\n    def clamp_weights(self):\n        for module in self.net.modules():\n            if hasattr(module, \"weight\") and module.kernel_size == (1, 1):\n                module.weight.data = torch.clamp(module.weight.data, min=0)\n\n    def set_input(self, data):\n        self.input_ref = data[\"ref\"]\n        self.input_p0 = data[\"p0\"]\n        self.input_p1 = data[\"p1\"]\n        self.input_judge = data[\"judge\"]\n\n        if self.use_gpu:\n            self.input_ref = self.input_ref.to(device=self.gpu_ids[0])\n            self.input_p0 = self.input_p0.to(device=self.gpu_ids[0])\n            self.input_p1 = self.input_p1.to(device=self.gpu_ids[0])\n            self.input_judge = self.input_judge.to(device=self.gpu_ids[0])\n\n        self.var_ref = Variable(self.input_ref, requires_grad=True)\n        self.var_p0 = Variable(self.input_p0, requires_grad=True)\n        self.var_p1 = Variable(self.input_p1, requires_grad=True)\n\n    def forward_train(self):  # run forward pass\n        # print(self.net.module.scaling_layer.shift)\n        # print(torch.norm(self.net.module.net.slice1[0].weight).item(), torch.norm(self.net.module.lin0.model[1].weight).item())\n\n        self.d0 = self.forward(self.var_ref, self.var_p0)\n        self.d1 = self.forward(self.var_ref, self.var_p1)\n        self.acc_r = self.compute_accuracy(self.d0, self.d1, self.input_judge)\n\n        self.var_judge = Variable(1.0 * self.input_judge).view(self.d0.size())\n\n        self.loss_total = self.rankLoss.forward(self.d0, self.d1, self.var_judge * 2.0 - 1.0)\n\n        return self.loss_total\n\n    def backward_train(self):\n        torch.mean(self.loss_total).backward()\n\n    def compute_accuracy(self, d0, d1, judge):\n        \"\"\" d0, d1 are Variables, judge is a Tensor \"\"\"\n        d1_lt_d0 = (d1 < d0).cpu().data.numpy().flatten()\n        judge_per = judge.cpu().numpy().flatten()\n        return d1_lt_d0 * judge_per + (1 - d1_lt_d0) * (1 - judge_per)\n\n    def get_current_errors(self):\n        retDict = OrderedDict([(\"loss_total\", self.loss_total.data.cpu().numpy()), (\"acc_r\", self.acc_r)])\n\n        for key in retDict.keys():\n            retDict[key] = np.mean(retDict[key])\n\n        return retDict\n\n    def get_current_visuals(self):\n        zoom_factor = 256 / self.var_ref.data.size()[2]\n\n        ref_img = util.tensor2im(self.var_ref.data)\n        p0_img = util.tensor2im(self.var_p0.data)\n        p1_img = util.tensor2im(self.var_p1.data)\n\n        ref_img_vis = zoom(ref_img, [zoom_factor, zoom_factor, 1], order=0)\n        p0_img_vis = zoom(p0_img, [zoom_factor, zoom_factor, 1], order=0)\n        p1_img_vis = zoom(p1_img, [zoom_factor, zoom_factor, 1], order=0)\n\n        return OrderedDict([(\"ref\", ref_img_vis), (\"p0\", p0_img_vis), (\"p1\", p1_img_vis)])\n\n    def save(self, path, label):\n        if self.use_gpu:\n            self.save_network(self.net.module, path, \"\", label)\n        else:\n            self.save_network(self.net, path, \"\", label)\n        self.save_network(self.rankLoss.net, path, \"rank\", label)\n\n    def update_learning_rate(self, nepoch_decay):\n        lrd = self.lr / nepoch_decay\n        lr = self.old_lr - lrd\n\n        for param_group in self.optimizer_net.param_groups:\n            param_group[\"lr\"] = lr\n\n        print(\"update lr [%s] decay: %f -> %f\" % (type, self.old_lr, lr))\n        self.old_lr = lr\n\n\ndef score_2afc_dataset(data_loader, func, name=\"\"):\n    \"\"\" Function computes Two Alternative Forced Choice (2AFC) score using\n        distance function 'func' in dataset 'data_loader'\n    INPUTS\n        data_loader - CustomDatasetDataLoader object - contains a TwoAFCDataset inside\n        func - callable distance function - calling d=func(in0,in1) should take 2\n            pytorch tensors with shape Nx3xXxY, and return numpy array of length N\n    OUTPUTS\n        [0] - 2AFC score in [0,1], fraction of time func agrees with human evaluators\n        [1] - dictionary with following elements\n            d0s,d1s - N arrays containing distances between reference patch to perturbed patches \n            gts - N array in [0,1], preferred patch selected by human evaluators\n                (closer to \"0\" for left patch p0, \"1\" for right patch p1,\n                \"0.6\" means 60pct people preferred right patch, 40pct preferred left)\n            scores - N array in [0,1], corresponding to what percentage function agreed with humans\n    CONSTS\n        N - number of test triplets in data_loader\n    \"\"\"\n\n    d0s = []\n    d1s = []\n    gts = []\n\n    for data in tqdm(data_loader.load_data(), desc=name):\n        d0s += func(data[\"ref\"], data[\"p0\"]).data.cpu().numpy().flatten().tolist()\n        d1s += func(data[\"ref\"], data[\"p1\"]).data.cpu().numpy().flatten().tolist()\n        gts += data[\"judge\"].cpu().numpy().flatten().tolist()\n\n    d0s = np.array(d0s)\n    d1s = np.array(d1s)\n    gts = np.array(gts)\n    scores = (d0s < d1s) * (1.0 - gts) + (d1s < d0s) * gts + (d1s == d0s) * 0.5\n\n    return (np.mean(scores), dict(d0s=d0s, d1s=d1s, gts=gts, scores=scores))\n\n\ndef score_jnd_dataset(data_loader, func, name=\"\"):\n    \"\"\" Function computes JND score using distance function 'func' in dataset 'data_loader'\n    INPUTS\n        data_loader - CustomDatasetDataLoader object - contains a JNDDataset inside\n        func - callable distance function - calling d=func(in0,in1) should take 2\n            pytorch tensors with shape Nx3xXxY, and return pytorch array of length N\n    OUTPUTS\n        [0] - JND score in [0,1], mAP score (area under precision-recall curve)\n        [1] - dictionary with following elements\n            ds - N array containing distances between two patches shown to human evaluator\n            sames - N array containing fraction of people who thought the two patches were identical\n    CONSTS\n        N - number of test triplets in data_loader\n    \"\"\"\n\n    ds = []\n    gts = []\n\n    for data in tqdm(data_loader.load_data(), desc=name):\n        ds += func(data[\"p0\"], data[\"p1\"]).data.cpu().numpy().tolist()\n        gts += data[\"same\"].cpu().numpy().flatten().tolist()\n\n    sames = np.array(gts)\n    ds = np.array(ds)\n\n    sorted_inds = np.argsort(ds)\n    ds_sorted = ds[sorted_inds]\n    sames_sorted = sames[sorted_inds]\n\n    TPs = np.cumsum(sames_sorted)\n    FPs = np.cumsum(1 - sames_sorted)\n    FNs = np.sum(sames_sorted) - TPs\n\n    precs = TPs / (TPs + FPs)\n    recs = TPs / (TPs + FNs)\n    score = util.voc_ap(recs, precs)\n\n    return (score, dict(ds=ds, sames=sames))\n\n"
  },
  {
    "path": "validation/lpips/networks_basic.py",
    "content": "import torch\nimport torch.nn as nn\nfrom torch.autograd import Variable\nfrom . import pretrained_networks as pn\n\nfrom . import util\n\n\ndef spatial_average(in_tens, keepdim=True):\n    return in_tens.mean([2, 3], keepdim=keepdim)\n\n\ndef upsample(in_tens, out_H=64):  # assumes scale factor is same for H and W\n    in_H = in_tens.shape[2]\n    scale_factor = 1.0 * out_H / in_H\n\n    return nn.Upsample(scale_factor=scale_factor, mode=\"bilinear\", align_corners=False)(in_tens)\n\n\n# Learned perceptual metric\nclass PNetLin(nn.Module):\n    def __init__(\n        self,\n        pnet_type=\"vgg\",\n        pnet_rand=False,\n        pnet_tune=False,\n        use_dropout=True,\n        spatial=False,\n        version=\"0.1\",\n        lpips=True,\n    ):\n        super(PNetLin, self).__init__()\n\n        self.pnet_type = pnet_type\n        self.pnet_tune = pnet_tune\n        self.pnet_rand = pnet_rand\n        self.spatial = spatial\n        self.lpips = lpips\n        self.version = version\n        self.scaling_layer = ScalingLayer()\n\n        if self.pnet_type in [\"vgg\", \"vgg16\"]:\n            net_type = pn.vgg16\n            self.chns = [64, 128, 256, 512, 512]\n        elif self.pnet_type == \"alex\":\n            net_type = pn.alexnet\n            self.chns = [64, 192, 384, 256, 256]\n        elif self.pnet_type == \"squeeze\":\n            net_type = pn.squeezenet\n            self.chns = [64, 128, 256, 384, 384, 512, 512]\n        self.L = len(self.chns)\n\n        self.net = net_type(pretrained=not self.pnet_rand, requires_grad=self.pnet_tune)\n\n        if lpips:\n            self.lin0 = NetLinLayer(self.chns[0], use_dropout=use_dropout)\n            self.lin1 = NetLinLayer(self.chns[1], use_dropout=use_dropout)\n            self.lin2 = NetLinLayer(self.chns[2], use_dropout=use_dropout)\n            self.lin3 = NetLinLayer(self.chns[3], use_dropout=use_dropout)\n            self.lin4 = NetLinLayer(self.chns[4], use_dropout=use_dropout)\n            self.lins = [self.lin0, self.lin1, self.lin2, self.lin3, self.lin4]\n            if self.pnet_type == \"squeeze\":  # 7 layers for squeezenet\n                self.lin5 = NetLinLayer(self.chns[5], use_dropout=use_dropout)\n                self.lin6 = NetLinLayer(self.chns[6], use_dropout=use_dropout)\n                self.lins += [self.lin5, self.lin6]\n\n    def forward(self, in0, in1, retPerLayer=False):\n        # v0.0 - original release had a bug, where input was not scaled\n        in0_input, in1_input = (\n            (self.scaling_layer(in0), self.scaling_layer(in1)) if self.version == \"0.1\" else (in0, in1)\n        )\n        outs0, outs1 = self.net.forward(in0_input), self.net.forward(in1_input)\n        feats0, feats1, diffs = {}, {}, {}\n\n        for kk in range(self.L):\n            feats0[kk], feats1[kk] = util.normalize_tensor(outs0[kk]), util.normalize_tensor(outs1[kk])\n            diffs[kk] = (feats0[kk] - feats1[kk]) ** 2\n\n        if self.lpips:\n            if self.spatial:\n                res = [upsample(self.lins[kk].model(diffs[kk]), out_H=in0.shape[2]) for kk in range(self.L)]\n            else:\n                res = [spatial_average(self.lins[kk].model(diffs[kk]), keepdim=True) for kk in range(self.L)]\n        else:\n            if self.spatial:\n                res = [upsample(diffs[kk].sum(dim=1, keepdim=True), out_H=in0.shape[2]) for kk in range(self.L)]\n            else:\n                res = [spatial_average(diffs[kk].sum(dim=1, keepdim=True), keepdim=True) for kk in range(self.L)]\n\n        val = res[0]\n        for l in range(1, self.L):\n            val += res[l]\n\n        if retPerLayer:\n            return (val, res)\n        else:\n            return val\n\n\nclass ScalingLayer(nn.Module):\n    def __init__(self):\n        super(ScalingLayer, self).__init__()\n        self.register_buffer(\"shift\", torch.Tensor([-0.030, -0.088, -0.188])[None, :, None, None])\n        self.register_buffer(\"scale\", torch.Tensor([0.458, 0.448, 0.450])[None, :, None, None])\n\n    def forward(self, inp):\n        return (inp - self.shift) / self.scale\n\n\nclass NetLinLayer(nn.Module):\n    \"\"\" A single linear layer which does a 1x1 conv \"\"\"\n\n    def __init__(self, chn_in, chn_out=1, use_dropout=False):\n        super(NetLinLayer, self).__init__()\n\n        layers = [nn.Dropout(),] if (use_dropout) else []\n        layers += [\n            nn.Conv2d(chn_in, chn_out, 1, stride=1, padding=0, bias=False),\n        ]\n        self.model = nn.Sequential(*layers)\n\n\nclass Dist2LogitLayer(nn.Module):\n    \"\"\" takes 2 distances, puts through fc layers, spits out value between [0,1] (if use_sigmoid is True) \"\"\"\n\n    def __init__(self, chn_mid=32, use_sigmoid=True):\n        super(Dist2LogitLayer, self).__init__()\n\n        layers = [\n            nn.Conv2d(5, chn_mid, 1, stride=1, padding=0, bias=True),\n        ]\n        layers += [\n            nn.LeakyReLU(0.2, True),\n        ]\n        layers += [\n            nn.Conv2d(chn_mid, chn_mid, 1, stride=1, padding=0, bias=True),\n        ]\n        layers += [\n            nn.LeakyReLU(0.2, True),\n        ]\n        layers += [\n            nn.Conv2d(chn_mid, 1, 1, stride=1, padding=0, bias=True),\n        ]\n        if use_sigmoid:\n            layers += [\n                nn.Sigmoid(),\n            ]\n        self.model = nn.Sequential(*layers)\n\n    def forward(self, d0, d1, eps=0.1):\n        return self.model.forward(torch.cat((d0, d1, d0 - d1, d0 / (d1 + eps), d1 / (d0 + eps)), dim=1))\n\n\nclass BCERankingLoss(nn.Module):\n    def __init__(self, chn_mid=32):\n        super(BCERankingLoss, self).__init__()\n        self.net = Dist2LogitLayer(chn_mid=chn_mid)\n        # self.parameters = list(self.net.parameters())\n        self.loss = torch.nn.BCELoss()\n\n    def forward(self, d0, d1, judge):\n        per = (judge + 1.0) / 2.0\n        self.logit = self.net.forward(d0, d1)\n        return self.loss(self.logit, per)\n\n\n# L2, DSSIM metrics\nclass FakeNet(nn.Module):\n    def __init__(self, use_gpu=True, colorspace=\"Lab\"):\n        super(FakeNet, self).__init__()\n        self.use_gpu = use_gpu\n        self.colorspace = colorspace\n\n\nclass L2(FakeNet):\n    def forward(self, in0, in1, retPerLayer=None):\n        assert in0.size()[0] == 1  # currently only supports batchSize 1\n\n        if self.colorspace == \"RGB\":\n            (N, C, X, Y) = in0.size()\n            value = torch.mean(\n                torch.mean(torch.mean((in0 - in1) ** 2, dim=1).view(N, 1, X, Y), dim=2).view(N, 1, 1, Y), dim=3\n            ).view(N)\n            return value\n        elif self.colorspace == \"Lab\":\n            value = util.l2(\n                util.tensor2np(util.tensor2tensorlab(in0.data, to_norm=False)),\n                util.tensor2np(util.tensor2tensorlab(in1.data, to_norm=False)),\n                range=100.0,\n            ).astype(\"float\")\n            ret_var = Variable(torch.Tensor((value,)))\n            if self.use_gpu:\n                ret_var = ret_var.cuda()\n            return ret_var\n\n\nclass DSSIM(FakeNet):\n    def forward(self, in0, in1, retPerLayer=None):\n        assert in0.size()[0] == 1  # currently only supports batchSize 1\n\n        if self.colorspace == \"RGB\":\n            value = util.dssim(1.0 * util.tensor2im(in0.data), 1.0 * util.tensor2im(in1.data), range=255.0).astype(\n                \"float\"\n            )\n        elif self.colorspace == \"Lab\":\n            value = util.dssim(\n                util.tensor2np(util.tensor2tensorlab(in0.data, to_norm=False)),\n                util.tensor2np(util.tensor2tensorlab(in1.data, to_norm=False)),\n                range=100.0,\n            ).astype(\"float\")\n        ret_var = Variable(torch.Tensor((value,)))\n        if self.use_gpu:\n            ret_var = ret_var.cuda()\n        return ret_var\n\n\ndef print_network(net):\n    num_params = 0\n    for param in net.parameters():\n        num_params += param.numel()\n    print(\"Network\", net)\n    print(\"Total number of parameters: %d\" % num_params)\n\n"
  },
  {
    "path": "validation/lpips/pretrained_networks.py",
    "content": "from collections import namedtuple\nimport torch\nfrom torchvision import models as tv\n\n\nclass squeezenet(torch.nn.Module):\n    def __init__(self, requires_grad=False, pretrained=True):\n        super(squeezenet, self).__init__()\n        pretrained_features = tv.squeezenet1_1(pretrained=pretrained).features\n        self.slice1 = torch.nn.Sequential()\n        self.slice2 = torch.nn.Sequential()\n        self.slice3 = torch.nn.Sequential()\n        self.slice4 = torch.nn.Sequential()\n        self.slice5 = torch.nn.Sequential()\n        self.slice6 = torch.nn.Sequential()\n        self.slice7 = torch.nn.Sequential()\n        self.N_slices = 7\n        for x in range(2):\n            self.slice1.add_module(str(x), pretrained_features[x])\n        for x in range(2, 5):\n            self.slice2.add_module(str(x), pretrained_features[x])\n        for x in range(5, 8):\n            self.slice3.add_module(str(x), pretrained_features[x])\n        for x in range(8, 10):\n            self.slice4.add_module(str(x), pretrained_features[x])\n        for x in range(10, 11):\n            self.slice5.add_module(str(x), pretrained_features[x])\n        for x in range(11, 12):\n            self.slice6.add_module(str(x), pretrained_features[x])\n        for x in range(12, 13):\n            self.slice7.add_module(str(x), pretrained_features[x])\n        if not requires_grad:\n            for param in self.parameters():\n                param.requires_grad = False\n\n    def forward(self, X):\n        h = self.slice1(X)\n        h_relu1 = h\n        h = self.slice2(h)\n        h_relu2 = h\n        h = self.slice3(h)\n        h_relu3 = h\n        h = self.slice4(h)\n        h_relu4 = h\n        h = self.slice5(h)\n        h_relu5 = h\n        h = self.slice6(h)\n        h_relu6 = h\n        h = self.slice7(h)\n        h_relu7 = h\n        vgg_outputs = namedtuple(\"SqueezeOutputs\", [\"relu1\", \"relu2\", \"relu3\", \"relu4\", \"relu5\", \"relu6\", \"relu7\"])\n        out = vgg_outputs(h_relu1, h_relu2, h_relu3, h_relu4, h_relu5, h_relu6, h_relu7)\n\n        return out\n\n\nclass alexnet(torch.nn.Module):\n    def __init__(self, requires_grad=False, pretrained=True):\n        super(alexnet, self).__init__()\n        alexnet_pretrained_features = tv.alexnet(pretrained=pretrained).features\n        self.slice1 = torch.nn.Sequential()\n        self.slice2 = torch.nn.Sequential()\n        self.slice3 = torch.nn.Sequential()\n        self.slice4 = torch.nn.Sequential()\n        self.slice5 = torch.nn.Sequential()\n        self.N_slices = 5\n        for x in range(2):\n            self.slice1.add_module(str(x), alexnet_pretrained_features[x])\n        for x in range(2, 5):\n            self.slice2.add_module(str(x), alexnet_pretrained_features[x])\n        for x in range(5, 8):\n            self.slice3.add_module(str(x), alexnet_pretrained_features[x])\n        for x in range(8, 10):\n            self.slice4.add_module(str(x), alexnet_pretrained_features[x])\n        for x in range(10, 12):\n            self.slice5.add_module(str(x), alexnet_pretrained_features[x])\n        if not requires_grad:\n            for param in self.parameters():\n                param.requires_grad = False\n\n    def forward(self, X):\n        h = self.slice1(X)\n        h_relu1 = h\n        h = self.slice2(h)\n        h_relu2 = h\n        h = self.slice3(h)\n        h_relu3 = h\n        h = self.slice4(h)\n        h_relu4 = h\n        h = self.slice5(h)\n        h_relu5 = h\n        alexnet_outputs = namedtuple(\"AlexnetOutputs\", [\"relu1\", \"relu2\", \"relu3\", \"relu4\", \"relu5\"])\n        out = alexnet_outputs(h_relu1, h_relu2, h_relu3, h_relu4, h_relu5)\n\n        return out\n\n\nclass vgg16(torch.nn.Module):\n    def __init__(self, requires_grad=False, pretrained=True):\n        super(vgg16, self).__init__()\n        vgg_pretrained_features = tv.vgg16(pretrained=pretrained).features\n        self.slice1 = torch.nn.Sequential()\n        self.slice2 = torch.nn.Sequential()\n        self.slice3 = torch.nn.Sequential()\n        self.slice4 = torch.nn.Sequential()\n        self.slice5 = torch.nn.Sequential()\n        self.N_slices = 5\n        for x in range(4):\n            self.slice1.add_module(str(x), vgg_pretrained_features[x])\n        for x in range(4, 9):\n            self.slice2.add_module(str(x), vgg_pretrained_features[x])\n        for x in range(9, 16):\n            self.slice3.add_module(str(x), vgg_pretrained_features[x])\n        for x in range(16, 23):\n            self.slice4.add_module(str(x), vgg_pretrained_features[x])\n        for x in range(23, 30):\n            self.slice5.add_module(str(x), vgg_pretrained_features[x])\n        if not requires_grad:\n            for param in self.parameters():\n                param.requires_grad = False\n\n    def forward(self, X):\n        h = self.slice1(X)\n        h_relu1_2 = h\n        h = self.slice2(h)\n        h_relu2_2 = h\n        h = self.slice3(h)\n        h_relu3_3 = h\n        h = self.slice4(h)\n        h_relu4_3 = h\n        h = self.slice5(h)\n        h_relu5_3 = h\n        vgg_outputs = namedtuple(\"VggOutputs\", [\"relu1_2\", \"relu2_2\", \"relu3_3\", \"relu4_3\", \"relu5_3\"])\n        out = vgg_outputs(h_relu1_2, h_relu2_2, h_relu3_3, h_relu4_3, h_relu5_3)\n\n        return out\n\n\nclass resnet(torch.nn.Module):\n    def __init__(self, requires_grad=False, pretrained=True, num=18):\n        super(resnet, self).__init__()\n        if num == 18:\n            self.net = tv.resnet18(pretrained=pretrained)\n        elif num == 34:\n            self.net = tv.resnet34(pretrained=pretrained)\n        elif num == 50:\n            self.net = tv.resnet50(pretrained=pretrained)\n        elif num == 101:\n            self.net = tv.resnet101(pretrained=pretrained)\n        elif num == 152:\n            self.net = tv.resnet152(pretrained=pretrained)\n        self.N_slices = 5\n\n        self.conv1 = self.net.conv1\n        self.bn1 = self.net.bn1\n        self.relu = self.net.relu\n        self.maxpool = self.net.maxpool\n        self.layer1 = self.net.layer1\n        self.layer2 = self.net.layer2\n        self.layer3 = self.net.layer3\n        self.layer4 = self.net.layer4\n\n    def forward(self, X):\n        h = self.conv1(X)\n        h = self.bn1(h)\n        h = self.relu(h)\n        h_relu1 = h\n        h = self.maxpool(h)\n        h = self.layer1(h)\n        h_conv2 = h\n        h = self.layer2(h)\n        h_conv3 = h\n        h = self.layer3(h)\n        h_conv4 = h\n        h = self.layer4(h)\n        h_conv5 = h\n\n        outputs = namedtuple(\"Outputs\", [\"relu1\", \"conv2\", \"conv3\", \"conv4\", \"conv5\"])\n        out = outputs(h_relu1, h_conv2, h_conv3, h_conv4, h_conv5)\n\n        return out\n"
  },
  {
    "path": "validation/lpips/util.py",
    "content": "from __future__ import absolute_import\nfrom __future__ import division\nfrom __future__ import print_function\n\nimport numpy as np\nfrom skimage.measure import compare_ssim\nimport torch\n\n\ndef normalize_tensor(in_feat, eps=1e-10):\n    norm_factor = torch.sqrt(torch.sum(in_feat ** 2, dim=1, keepdim=True))\n    return in_feat / (norm_factor + eps)\n\n\ndef l2(p0, p1, range=255.0):\n    return 0.5 * np.mean((p0 / range - p1 / range) ** 2)\n\n\ndef psnr(p0, p1, peak=255.0):\n    return 10 * np.log10(peak ** 2 / np.mean((1.0 * p0 - 1.0 * p1) ** 2))\n\n\ndef dssim(p0, p1, range=255.0):\n    return (1 - compare_ssim(p0, p1, data_range=range, multichannel=True)) / 2.0\n\n\ndef rgb2lab(in_img, mean_cent=False):\n    from skimage import color\n\n    img_lab = color.rgb2lab(in_img)\n    if mean_cent:\n        img_lab[:, :, 0] = img_lab[:, :, 0] - 50\n    return img_lab\n\n\ndef tensor2np(tensor_obj):\n    # change dimension of a tensor object into a numpy array\n    return tensor_obj[0].cpu().float().numpy().transpose((1, 2, 0))\n\n\ndef np2tensor(np_obj):\n    # change dimenion of np array into tensor array\n    return torch.Tensor(np_obj[:, :, :, np.newaxis].transpose((3, 2, 0, 1)))\n\n\ndef tensor2tensorlab(image_tensor, to_norm=True, mc_only=False):\n    # image tensor to lab tensor\n    from skimage import color\n\n    img = tensor2im(image_tensor)\n    img_lab = color.rgb2lab(img)\n    if mc_only:\n        img_lab[:, :, 0] = img_lab[:, :, 0] - 50\n    if to_norm and not mc_only:\n        img_lab[:, :, 0] = img_lab[:, :, 0] - 50\n        img_lab = img_lab / 100.0\n\n    return np2tensor(img_lab)\n\n\ndef tensorlab2tensor(lab_tensor, return_inbnd=False):\n    from skimage import color\n    import warnings\n\n    warnings.filterwarnings(\"ignore\")\n\n    lab = tensor2np(lab_tensor) * 100.0\n    lab[:, :, 0] = lab[:, :, 0] + 50\n\n    rgb_back = 255.0 * np.clip(color.lab2rgb(lab.astype(\"float\")), 0, 1)\n    if return_inbnd:\n        # convert back to lab, see if we match\n        lab_back = color.rgb2lab(rgb_back.astype(\"uint8\"))\n        mask = 1.0 * np.isclose(lab_back, lab, atol=2.0)\n        mask = np2tensor(np.prod(mask, axis=2)[:, :, np.newaxis])\n        return (im2tensor(rgb_back), mask)\n    else:\n        return im2tensor(rgb_back)\n\n\ndef rgb2lab(input):\n    from skimage import color\n\n    return color.rgb2lab(input / 255.0)\n\n\ndef tensor2im(image_tensor, imtype=np.uint8, cent=1.0, factor=255.0 / 2.0):\n    image_numpy = image_tensor[0].cpu().float().numpy()\n    image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + cent) * factor\n    return image_numpy.astype(imtype)\n\n\ndef im2tensor(image, imtype=np.uint8, cent=1.0, factor=255.0 / 2.0):\n    return torch.Tensor((image / factor - cent)[:, :, :, np.newaxis].transpose((3, 2, 0, 1)))\n\n\ndef tensor2vec(vector_tensor):\n    return vector_tensor.data.cpu().numpy()[:, :, 0, 0]\n\n\ndef voc_ap(rec, prec, use_07_metric=False):\n    \"\"\" ap = voc_ap(rec, prec, [use_07_metric])\n    Compute VOC AP given precision and recall.\n    If use_07_metric is true, uses the\n    VOC 07 11 point method (default:False).\n    \"\"\"\n    if use_07_metric:\n        # 11 point metric\n        ap = 0.0\n        for t in np.arange(0.0, 1.1, 0.1):\n            if np.sum(rec >= t) == 0:\n                p = 0\n            else:\n                p = np.max(prec[rec >= t])\n            ap = ap + p / 11.0\n    else:\n        # correct AP calculation\n        # first append sentinel values at the end\n        mrec = np.concatenate(([0.0], rec, [1.0]))\n        mpre = np.concatenate(([0.0], prec, [0.0]))\n\n        # compute the precision envelope\n        for i in range(mpre.size - 1, 0, -1):\n            mpre[i - 1] = np.maximum(mpre[i - 1], mpre[i])\n\n        # to calculate area under PR curve, look for points\n        # where X axis (recall) changes value\n        i = np.where(mrec[1:] != mrec[:-1])[0]\n\n        # and sum (\\Delta recall) * prec\n        ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1])\n    return ap\n\n\ndef tensor2im(image_tensor, imtype=np.uint8, cent=1.0, factor=255.0 / 2.0):\n    # def tensor2im(image_tensor, imtype=np.uint8, cent=1., factor=1.):\n    image_numpy = image_tensor[0].cpu().float().numpy()\n    image_numpy = (np.transpose(image_numpy, (1, 2, 0)) + cent) * factor\n    return image_numpy.astype(imtype)\n\n\ndef im2tensor(image, imtype=np.uint8, cent=1.0, factor=255.0 / 2.0):\n    # def im2tensor(image, imtype=np.uint8, cent=1., factor=1.):\n    return torch.Tensor((image / factor - cent)[:, :, :, np.newaxis].transpose((3, 2, 0, 1)))\n\n"
  },
  {
    "path": "validation/metrics.py",
    "content": "import os\nimport pickle\nimport random\n\nfrom sklearn.metrics import pairwise_distances\nfrom tqdm import tqdm\nimport torch\nfrom torch.nn import functional as F\nimport numpy as np\nfrom scipy import linalg\n\nfrom .inception import InceptionV3\nfrom . import lpips\n\n\n@torch.no_grad()\ndef vae_fid(vae, batch_size, latent_dim, n_sample, inception_name, calculate_prdc=True):\n    vae.eval()\n\n    inception = InceptionV3([3], normalize_input=False, init_weights=False)\n    inception = inception.eval().to(next(vae.parameters()).device)\n\n    n_batch = n_sample // batch_size\n    resid = n_sample - (n_batch * batch_size)\n    if resid == 0:\n        batch_sizes = [batch_size] * n_batch\n    else:\n        batch_sizes = [batch_size] * n_batch + [resid]\n    features = []\n\n    for batch in batch_sizes:\n        latent = torch.randn(batch, *latent_dim).cuda()\n        img = vae.decode(latent)\n        feat = inception(img)[0].view(img.shape[0], -1)\n        features.append(feat.to(\"cpu\"))\n    features = torch.cat(features, 0).numpy()\n\n    del inception\n\n    sample_mean = np.mean(features, 0)\n    sample_cov = np.cov(features, rowvar=False)\n\n    with open(f\"inception_{inception_name}_stats.pkl\", \"rb\") as f:\n        embeds = pickle.load(f)\n        real_mean = embeds[\"mean\"]\n        real_cov = embeds[\"cov\"]\n\n    cov_sqrt, _ = linalg.sqrtm(sample_cov @ real_cov, disp=False)\n\n    if not np.isfinite(cov_sqrt).all():\n        print(\"product of cov matrices is singular\")\n        offset = np.eye(sample_cov.shape[0]) * 1e-6\n        cov_sqrt = linalg.sqrtm((sample_cov + offset) @ (real_cov + offset))\n\n    if np.iscomplexobj(cov_sqrt):\n        if not np.allclose(np.diagonal(cov_sqrt).imag, 0, atol=1e-3):\n            m = np.max(np.abs(cov_sqrt.imag))\n\n            raise ValueError(f\"Imaginary component {m}\")\n\n        cov_sqrt = cov_sqrt.real\n\n    mean_diff = sample_mean - real_mean\n    mean_norm = mean_diff @ mean_diff\n\n    trace = np.trace(sample_cov) + np.trace(real_cov) - 2 * np.trace(cov_sqrt)\n\n    fid = mean_norm + trace\n\n    ret_dict = {\"FID\": fid}\n\n    if calculate_prdc:\n        with open(f\"inception_{inception_name}_features.pkl\", \"rb\") as f:\n            embeds = pickle.load(f)\n            real_feats = embeds[\"features\"]\n        _, _, density, coverage = prdc(real_feats[:80000], features[:80000])\n        ret_dict[\"Density\"] = density\n        ret_dict[\"Coverage\"] = coverage\n\n    return ret_dict\n\n\n@torch.no_grad()\ndef fid(generator, batch_size, n_sample, truncation, inception_name, calculate_prdc=True):\n    generator.eval()\n    mean_latent = generator.mean_latent(2 ** 14)\n\n    inception = InceptionV3([3], normalize_input=False, init_weights=False)\n    inception = inception.eval().to(next(generator.parameters()).device)\n\n    n_batch = n_sample // batch_size\n    resid = n_sample - (n_batch * batch_size)\n    if resid == 0:\n        batch_sizes = [batch_size] * n_batch\n    else:\n        batch_sizes = [batch_size] * n_batch + [resid]\n    features = []\n\n    for batch in batch_sizes:\n        if truncation is None:\n            trunc = random.uniform(0.9, 1.5)\n        else:\n            trunc = truncation\n        latent = torch.randn(batch, 512).cuda()\n        img, _ = generator([latent], truncation=trunc, truncation_latent=mean_latent)\n        feat = inception(img)[0].view(img.shape[0], -1)\n        features.append(feat.to(\"cpu\"))\n    features = torch.cat(features, 0).numpy()\n\n    del inception\n\n    sample_mean = np.mean(features, 0)\n    sample_cov = np.cov(features, rowvar=False)\n\n    with open(f\"inception_{inception_name}_stats.pkl\", \"rb\") as f:\n        embeds = pickle.load(f)\n        real_mean = embeds[\"mean\"]\n        real_cov = embeds[\"cov\"]\n\n    cov_sqrt, _ = linalg.sqrtm(sample_cov @ real_cov, disp=False)\n\n    if not np.isfinite(cov_sqrt).all():\n        print(\"product of cov matrices is singular\")\n        offset = np.eye(sample_cov.shape[0]) * 1e-6\n        cov_sqrt = linalg.sqrtm((sample_cov + offset) @ (real_cov + offset))\n\n    if np.iscomplexobj(cov_sqrt):\n        if not np.allclose(np.diagonal(cov_sqrt).imag, 0, atol=1e-3):\n            m = np.max(np.abs(cov_sqrt.imag))\n\n            raise ValueError(f\"Imaginary component {m}\")\n\n        cov_sqrt = cov_sqrt.real\n\n    mean_diff = sample_mean - real_mean\n    mean_norm = mean_diff @ mean_diff\n\n    trace = np.trace(sample_cov) + np.trace(real_cov) - 2 * np.trace(cov_sqrt)\n\n    fid = mean_norm + trace\n\n    ret_dict = {\"FID\": fid}\n\n    if calculate_prdc:\n        with open(f\"inception_{inception_name}_features.pkl\", \"rb\") as f:\n            embeds = pickle.load(f)\n            real_feats = embeds[\"features\"]\n        _, _, density, coverage = prdc(real_feats[:80000], features[:80000])\n        ret_dict[\"Density\"] = density\n        ret_dict[\"Coverage\"] = coverage\n\n    return ret_dict\n\n\ndef get_dataset_inception_features(loader, inception_name, size):\n    if not os.path.exists(f\"inception_{inception_name}_stats.pkl\"):\n        print(\"calculating inception features for FID....\")\n        inception = InceptionV3([3], normalize_input=False, init_weights=False)\n        inception = torch.nn.DataParallel(inception).eval().cuda()\n\n        feature_list = []\n        for img in tqdm(loader):\n            img = img.cuda()\n            feature = inception(img)[0].view(img.shape[0], -1)\n            feature_list.append(feature.to(\"cpu\"))\n        features = torch.cat(feature_list, 0).numpy()\n\n        mean = np.mean(features, 0)\n        cov = np.cov(features, rowvar=False)\n\n        with open(f\"inception_{inception_name}_stats.pkl\", \"wb\") as f:\n            pickle.dump({\"mean\": mean, \"cov\": cov, \"size\": size, \"feat\": features}, f)\n        with open(f\"inception_{inception_name}_features.pkl\", \"wb\") as f:\n            pickle.dump({\"features\": features}, f)\n    else:\n        print(f\"Found inception features: inception_{inception_name}_stats.pkl\")\n\n\ndef compute_pairwise_distance(data_x, data_y=None, metric=\"l2\"):\n    if data_y is None:\n        data_y = data_x\n    dists = pairwise_distances(\n        data_x.reshape((len(data_x), -1)), data_y.reshape((len(data_y), -1)), metric=metric, n_jobs=24\n    )\n    return dists\n\n\ndef get_kth_value(unsorted, k, axis=-1):\n    indices = np.argpartition(unsorted, k, axis=axis)[..., :k]\n    k_smallests = np.take_along_axis(unsorted, indices, axis=axis)\n    kth_values = k_smallests.max(axis=axis)\n    return kth_values\n\n\ndef compute_nearest_neighbour_distances(input_features, nearest_k, metric):\n    distances = compute_pairwise_distance(input_features, metric=metric)\n    radii = get_kth_value(distances, k=nearest_k + 1, axis=-1)\n    return radii\n\n\ndef prdc(real_features, fake_features, nearest_k=10, metric=\"l2\"):\n    real_nearest_neighbour_distances = compute_nearest_neighbour_distances(real_features, nearest_k, metric=metric)\n    fake_nearest_neighbour_distances = compute_nearest_neighbour_distances(fake_features, nearest_k, metric=metric)\n    distance_real_fake = compute_pairwise_distance(real_features, fake_features, metric=metric)\n\n    precision = (distance_real_fake < np.expand_dims(real_nearest_neighbour_distances, axis=1)).any(axis=0).mean()\n    recall = (distance_real_fake < np.expand_dims(fake_nearest_neighbour_distances, axis=0)).any(axis=1).mean()\n\n    density = (1.0 / float(nearest_k)) * (\n        distance_real_fake < np.expand_dims(real_nearest_neighbour_distances, axis=1)\n    ).sum(axis=0).mean()\n    coverage = (distance_real_fake.min(axis=1) < real_nearest_neighbour_distances).mean()\n\n    return precision, recall, density, coverage\n\n\ndef lerp(a, b, t):\n    return a + (b - a) * t\n\n\n@torch.no_grad()\ndef ppl(generator, batch_size, n_sample, space, crop, latent_dim, eps=1e-4):\n    generator.eval()\n\n    percept = lpips.PerceptualLoss(\n        model=\"net-lin\", net=\"vgg\", use_gpu=True, gpu_ids=[next(generator.parameters()).device.index]\n    )\n\n    distances = []\n\n    n_batch = n_sample // batch_size\n    resid = n_sample - (n_batch * batch_size)\n    if resid == 0:\n        batch_sizes = [batch_size] * n_batch\n    else:\n        batch_sizes = [batch_size] * n_batch + [resid]\n\n    for batch_size in batch_sizes:\n        noise = generator.make_noise()\n\n        inputs = torch.randn([batch_size * 2, latent_dim]).cuda()\n        lerp_t = torch.rand(batch_size).cuda()\n\n        if space == \"w\":\n            latent = generator.get_latent(inputs)\n            latent_t0, latent_t1 = latent[::2], latent[1::2]\n            latent_e0 = lerp(latent_t0, latent_t1, lerp_t[:, None])\n            latent_e1 = lerp(latent_t0, latent_t1, lerp_t[:, None] + eps)\n            latent_e = torch.stack([latent_e0, latent_e1], 1).view(*latent.shape)\n\n        image, _ = generator(latent_e, input_is_latent=True, noise=noise)\n\n        if crop:\n            c = image.shape[2] // 8\n            image = image[:, :, c * 3 : c * 7, c * 2 : c * 6]\n\n        factor = image.shape[2] // 256\n\n        if factor > 1:\n            image = F.interpolate(image, size=(256, 256), mode=\"bilinear\", align_corners=False)\n\n        dist = percept(image[::2], image[1::2]).view(image.shape[0] // 2) / (eps ** 2)\n        distances.append(dist.to(\"cpu\").numpy())\n\n    distances = np.concatenate(distances, 0)\n\n    lo = np.percentile(distances, 1, interpolation=\"lower\")\n    hi = np.percentile(distances, 99, interpolation=\"higher\")\n    filtered_dist = np.extract(np.logical_and(lo <= distances, distances <= hi), distances)\n    path_length = filtered_dist.mean()\n\n    del percept, inputs, lerp_t, image, dist\n\n    return path_length\n"
  },
  {
    "path": "validation/spectral_norm.py",
    "content": "import torch\n\n\nclass SpectralNorm(object):\n    def __init__(self, name=\"weight\", n_power_iterations=1, dim=0, eps=1e-12):\n        self.name = name\n        self.dim = dim\n        if n_power_iterations <= 0:\n            raise ValueError(\n                \"Expected n_power_iterations to be positive, but \"\n                \"got n_power_iterations={}\".format(n_power_iterations)\n            )\n        self.n_power_iterations = n_power_iterations\n        self.eps = eps\n\n    def reshape_weight_to_matrix(self, weight):\n        weight_mat = weight\n        if self.dim != 0:\n            # permute dim to front\n            weight_mat = weight_mat.permute(self.dim, *[d for d in range(weight_mat.dim()) if d != self.dim])\n        height = weight_mat.size(0)\n        return weight_mat.reshape(height, -1)\n\n    def compute_sigma(self, module):\n        with torch.no_grad():\n            weight = getattr(module, self.name)\n            weight_mat = self.reshape_weight_to_matrix(weight)\n\n            u = getattr(module, self.name + \"_u\")\n            v = getattr(module, self.name + \"_v\")\n            for _ in range(self.n_power_iterations):\n                v = torch.nn.functional.normalize(torch.mv(weight_mat.t(), u), dim=0, eps=self.eps)\n                u = torch.nn.functional.normalize(torch.mv(weight_mat, v), dim=0, eps=self.eps)\n            setattr(module, self.name + \"_u\", u)\n            setattr(module, self.name + \"_v\", v)\n\n            sigma = torch.dot(u, torch.mv(weight_mat, v))\n            setattr(module, \"spectral_norm\", sigma)\n\n    def remove(self, module):\n        delattr(module, self.name)\n        delattr(module, self.name + \"_u\")\n        delattr(module, self.name + \"_v\")\n        delattr(module, \"spectral_norm\")\n\n    def __call__(self, module, inputs):\n        self.compute_sigma(module)\n\n    def _solve_v_and_rescale(self, weight_mat, u, target_sigma):\n        v = torch.chain_matmul(weight_mat.t().mm(weight_mat).pinverse(), weight_mat.t(), u.unsqueeze(1)).squeeze(1)\n        return v.mul_(target_sigma / torch.dot(u, torch.mv(weight_mat, v)))\n\n    @staticmethod\n    def apply(module, name, n_power_iterations, dim, eps, normalize=True):\n        for k, hook in module._forward_pre_hooks.items():\n            if isinstance(hook, SpectralNorm) and hook.name == name:\n                raise RuntimeError(\"Cannot register two spectral_norm hooks on \" \"the same parameter {}\".format(name))\n\n        fn = SpectralNorm(name, n_power_iterations, dim, eps)\n        weight = module._parameters[name]\n        with torch.no_grad():\n            weight_mat = fn.reshape_weight_to_matrix(weight)\n\n            h, w = weight_mat.size()\n            u = torch.nn.functional.normalize(weight.new_empty(h).normal_(0, 1), dim=0, eps=fn.eps)\n            v = torch.nn.functional.normalize(weight.new_empty(w).normal_(0, 1), dim=0, eps=fn.eps)\n\n        module.register_buffer(fn.name + \"_u\", u)\n        module.register_buffer(fn.name + \"_v\", v)\n        module.register_buffer(\"spectral_norm\", torch.tensor(-1, device=next(module.parameters()).device))\n\n        module.register_forward_pre_hook(fn)\n        return fn\n\n\ndef track_spectral_norm(module, name=\"weight\", n_power_iterations=1, eps=1e-12, dim=None):\n    r\"\"\"Tracks the spectral norm of a module's weight parameter\n    Args:\n        module (nn.Module): containing module\n        name (str, optional): name of weight parameter\n        n_power_iterations (int, optional): number of power iterations to\n            calculate spectral norm\n        eps (float, optional): epsilon for numerical stability in\n            calculating norms\n        dim (int, optional): dimension corresponding to number of outputs,\n            the default is ``0``, except for modules that are instances of\n            ConvTranspose{1,2,3}d, when it is ``1``\n    Returns:\n        The original module with the spectral norm hook\n    Example::\n        >>> m = spectral_norm(nn.Linear(20, 40))\n        >>> m\n        Linear(in_features=20, out_features=40, bias=True)\n        >>> m.weight_u.size()\n        torch.Size([40])\n    \"\"\"\n    if dim is None:\n        if isinstance(module, (torch.nn.ConvTranspose1d, torch.nn.ConvTranspose2d, torch.nn.ConvTranspose3d)):\n            dim = 1\n        else:\n            dim = 0\n    SpectralNorm.apply(module, name, n_power_iterations, dim, eps)\n    return module\n\n\ndef remove_spectral_norm(module, name=\"weight\"):\n    r\"\"\"Removes the spectral normalization reparameterization from a module.\n    Args:\n        module (Module): containing module\n        name (str, optional): name of weight parameter\n    Example:\n        >>> m = spectral_norm(nn.Linear(40, 10))\n        >>> remove_spectral_norm(m)\n    \"\"\"\n    for k, hook in module._forward_pre_hooks.items():\n        if isinstance(hook, SpectralNorm) and hook.name == name:\n            hook.remove(module)\n            del module._forward_pre_hooks[k]\n            break\n    else:\n        raise ValueError(\"spectral_norm of '{}' not found in {}\".format(name, module))\n\n    return module\n"
  },
  {
    "path": "workspace/naamloos_metadata.json",
    "content": "{\"total_frames\": 4986}\n"
  },
  {
    "path": "workspace/naamloos_params.json",
    "content": "{\"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}"
  }
]