Training in progress, step 1000

.gitattributes

@@ -33,3 +33,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+huggan/assets/cyclegan.png filter=lfs diff=lfs merge=lfs -text
+huggan/assets/lightweight_gan_wandb.png filter=lfs diff=lfs merge=lfs -text
+huggan/assets/pix2pix_maps.png filter=lfs diff=lfs merge=lfs -text
+huggan/assets/wandb.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,166 @@
+# Initially taken from Github's Python gitignore file
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# tests and logs
+tests/fixtures/cached_*_text.txt
+logs/
+lightning_logs/
+lang_code_data/
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# vscode
+.vs
+.vscode
+
+# Pycharm
+.idea
+
+# TF code
+tensorflow_code
+
+# Models
+proc_data
+
+# examples
+runs
+/runs_old
+/wandb
+/examples/runs
+/examples/**/*.args
+/examples/rag/sweep
+
+# data
+/data
+serialization_dir
+
+# emacs
+*.*~
+debug.env
+
+# vim
+.*.swp
+
+#ctags
+tags
+
+# pre-commit
+.pre-commit*
+
+# .lock
+*.lock
+
+# DS_Store (MacOS)
+.DS_Store

README.md

@@ -0,0 +1,16 @@
+# Community Events @ 🤗
+
+A central repository for all community events organized by 🤗 HuggingFace. Come one, come all!
+We're constantly finding ways to democratise the use of ML across modalities and languages. This repo contains information about all past, present and upcoming events.
+
+## Hugging Events
+
+| **Event Name**                                                          | **Dates**       | **Status**                                                                                                   |
+|-------------------------------------------------------------------------|-----------------|--------------------------------------------------------------------------------------------------------------|
+| [Open Source AI Game Jam 🎮 (First Edition)](/open-source-ai-game-jam)               | July 7th - 9th, 2023  | Finished |
+| [Whisper Fine Tuning Event](/whisper-fine-tuning-event)               | Dec 5th - 19th, 2022  | Finished |
+| [Computer Vision Study Group](/computer-vision-study-group)             | Ongoing         | Monthly                                                                                                      |
+| [ML for Audio Study Group](https://github.com/Vaibhavs10/ml-with-audio) | Ongoing         | Monthly                                                                                                      |
+| [Gradio Blocks](/gradio-blocks)                                         | May 16th - 31st, 2022 | Finished                                                                                                     |
+| [HugGAN](/huggan)                                                       | Apr 4th - 17th, 2022  | Finished                                                                                                     |
+| [Keras Sprint](keras-sprint)                                            | June, 2022            | Finished                                                                                                     |

added_tokens.json

@@ -0,0 +1,1609 @@
+{
+  "<|0.00|>": 50364,
+  "<|0.02|>": 50365,
+  "<|0.04|>": 50366,
+  "<|0.06|>": 50367,
+  "<|0.08|>": 50368,
+  "<|0.10|>": 50369,
+  "<|0.12|>": 50370,
+  "<|0.14|>": 50371,
+  "<|0.16|>": 50372,
+  "<|0.18|>": 50373,
+  "<|0.20|>": 50374,
+  "<|0.22|>": 50375,
+  "<|0.24|>": 50376,
+  "<|0.26|>": 50377,
+  "<|0.28|>": 50378,
+  "<|0.30|>": 50379,
+  "<|0.32|>": 50380,
+  "<|0.34|>": 50381,
+  "<|0.36|>": 50382,
+  "<|0.38|>": 50383,
+  "<|0.40|>": 50384,
+  "<|0.42|>": 50385,
+  "<|0.44|>": 50386,
+  "<|0.46|>": 50387,
+  "<|0.48|>": 50388,
+  "<|0.50|>": 50389,
+  "<|0.52|>": 50390,
+  "<|0.54|>": 50391,
+  "<|0.56|>": 50392,
+  "<|0.58|>": 50393,
+  "<|0.60|>": 50394,
+  "<|0.62|>": 50395,
+  "<|0.64|>": 50396,
+  "<|0.66|>": 50397,
+  "<|0.68|>": 50398,
+  "<|0.70|>": 50399,
+  "<|0.72|>": 50400,
+  "<|0.74|>": 50401,
+  "<|0.76|>": 50402,
+  "<|0.78|>": 50403,
+  "<|0.80|>": 50404,
+  "<|0.82|>": 50405,
+  "<|0.84|>": 50406,
+  "<|0.86|>": 50407,
+  "<|0.88|>": 50408,
+  "<|0.90|>": 50409,
+  "<|0.92|>": 50410,
+  "<|0.94|>": 50411,
+  "<|0.96|>": 50412,
+  "<|0.98|>": 50413,
+  "<|1.00|>": 50414,
+  "<|1.02|>": 50415,
+  "<|1.04|>": 50416,
+  "<|1.06|>": 50417,
+  "<|1.08|>": 50418,
+  "<|1.10|>": 50419,
+  "<|1.12|>": 50420,
+  "<|1.14|>": 50421,
+  "<|1.16|>": 50422,
+  "<|1.18|>": 50423,
+  "<|1.20|>": 50424,
+  "<|1.22|>": 50425,
+  "<|1.24|>": 50426,
+  "<|1.26|>": 50427,
+  "<|1.28|>": 50428,
+  "<|1.30|>": 50429,
+  "<|1.32|>": 50430,
+  "<|1.34|>": 50431,
+  "<|1.36|>": 50432,
+  "<|1.38|>": 50433,
+  "<|1.40|>": 50434,
+  "<|1.42|>": 50435,
+  "<|1.44|>": 50436,
+  "<|1.46|>": 50437,
+  "<|1.48|>": 50438,
+  "<|1.50|>": 50439,
+  "<|1.52|>": 50440,
+  "<|1.54|>": 50441,
+  "<|1.56|>": 50442,
+  "<|1.58|>": 50443,
+  "<|1.60|>": 50444,
+  "<|1.62|>": 50445,
+  "<|1.64|>": 50446,
+  "<|1.66|>": 50447,
+  "<|1.68|>": 50448,
+  "<|1.70|>": 50449,
+  "<|1.72|>": 50450,
+  "<|1.74|>": 50451,
+  "<|1.76|>": 50452,
+  "<|1.78|>": 50453,
+  "<|1.80|>": 50454,
+  "<|1.82|>": 50455,
+  "<|1.84|>": 50456,
+  "<|1.86|>": 50457,
+  "<|1.88|>": 50458,
+  "<|1.90|>": 50459,
+  "<|1.92|>": 50460,
+  "<|1.94|>": 50461,
+  "<|1.96|>": 50462,
+  "<|1.98|>": 50463,
+  "<|10.00|>": 50864,
+  "<|10.02|>": 50865,
+  "<|10.04|>": 50866,
+  "<|10.06|>": 50867,
+  "<|10.08|>": 50868,
+  "<|10.10|>": 50869,
+  "<|10.12|>": 50870,
+  "<|10.14|>": 50871,
+  "<|10.16|>": 50872,
+  "<|10.18|>": 50873,
+  "<|10.20|>": 50874,
+  "<|10.22|>": 50875,
+  "<|10.24|>": 50876,
+  "<|10.26|>": 50877,
+  "<|10.28|>": 50878,
+  "<|10.30|>": 50879,
+  "<|10.32|>": 50880,
+  "<|10.34|>": 50881,
+  "<|10.36|>": 50882,
+  "<|10.38|>": 50883,
+  "<|10.40|>": 50884,
+  "<|10.42|>": 50885,
+  "<|10.44|>": 50886,
+  "<|10.46|>": 50887,
+  "<|10.48|>": 50888,
+  "<|10.50|>": 50889,
+  "<|10.52|>": 50890,
+  "<|10.54|>": 50891,
+  "<|10.56|>": 50892,
+  "<|10.58|>": 50893,
+  "<|10.60|>": 50894,
+  "<|10.62|>": 50895,
+  "<|10.64|>": 50896,
+  "<|10.66|>": 50897,
+  "<|10.68|>": 50898,
+  "<|10.70|>": 50899,
+  "<|10.72|>": 50900,
+  "<|10.74|>": 50901,
+  "<|10.76|>": 50902,
+  "<|10.78|>": 50903,
+  "<|10.80|>": 50904,
+  "<|10.82|>": 50905,
+  "<|10.84|>": 50906,
+  "<|10.86|>": 50907,
+  "<|10.88|>": 50908,
+  "<|10.90|>": 50909,
+  "<|10.92|>": 50910,
+  "<|10.94|>": 50911,
+  "<|10.96|>": 50912,
+  "<|10.98|>": 50913,
+  "<|11.00|>": 50914,
+  "<|11.02|>": 50915,
+  "<|11.04|>": 50916,
+  "<|11.06|>": 50917,
+  "<|11.08|>": 50918,
+  "<|11.10|>": 50919,
+  "<|11.12|>": 50920,
+  "<|11.14|>": 50921,
+  "<|11.16|>": 50922,
+  "<|11.18|>": 50923,
+  "<|11.20|>": 50924,
+  "<|11.22|>": 50925,
+  "<|11.24|>": 50926,
+  "<|11.26|>": 50927,
+  "<|11.28|>": 50928,
+  "<|11.30|>": 50929,
+  "<|11.32|>": 50930,
+  "<|11.34|>": 50931,
+  "<|11.36|>": 50932,
+  "<|11.38|>": 50933,
+  "<|11.40|>": 50934,
+  "<|11.42|>": 50935,
+  "<|11.44|>": 50936,
+  "<|11.46|>": 50937,
+  "<|11.48|>": 50938,
+  "<|11.50|>": 50939,
+  "<|11.52|>": 50940,
+  "<|11.54|>": 50941,
+  "<|11.56|>": 50942,
+  "<|11.58|>": 50943,
+  "<|11.60|>": 50944,
+  "<|11.62|>": 50945,
+  "<|11.64|>": 50946,
+  "<|11.66|>": 50947,
+  "<|11.68|>": 50948,
+  "<|11.70|>": 50949,
+  "<|11.72|>": 50950,
+  "<|11.74|>": 50951,
+  "<|11.76|>": 50952,
+  "<|11.78|>": 50953,
+  "<|11.80|>": 50954,
+  "<|11.82|>": 50955,
+  "<|11.84|>": 50956,
+  "<|11.86|>": 50957,
+  "<|11.88|>": 50958,
+  "<|11.90|>": 50959,
+  "<|11.92|>": 50960,
+  "<|11.94|>": 50961,
+  "<|11.96|>": 50962,
+  "<|11.98|>": 50963,
+  "<|12.00|>": 50964,
+  "<|12.02|>": 50965,
+  "<|12.04|>": 50966,
+  "<|12.06|>": 50967,
+  "<|12.08|>": 50968,
+  "<|12.10|>": 50969,
+  "<|12.12|>": 50970,
+  "<|12.14|>": 50971,
+  "<|12.16|>": 50972,
+  "<|12.18|>": 50973,
+  "<|12.20|>": 50974,
+  "<|12.22|>": 50975,
+  "<|12.24|>": 50976,
+  "<|12.26|>": 50977,
+  "<|12.28|>": 50978,
+  "<|12.30|>": 50979,
+  "<|12.32|>": 50980,
+  "<|12.34|>": 50981,
+  "<|12.36|>": 50982,
+  "<|12.38|>": 50983,
+  "<|12.40|>": 50984,
+  "<|12.42|>": 50985,
+  "<|12.44|>": 50986,
+  "<|12.46|>": 50987,
+  "<|12.48|>": 50988,
+  "<|12.50|>": 50989,
+  "<|12.52|>": 50990,
+  "<|12.54|>": 50991,
+  "<|12.56|>": 50992,
+  "<|12.58|>": 50993,
+  "<|12.60|>": 50994,
+  "<|12.62|>": 50995,
+  "<|12.64|>": 50996,
+  "<|12.66|>": 50997,
+  "<|12.68|>": 50998,
+  "<|12.70|>": 50999,
+  "<|12.72|>": 51000,
+  "<|12.74|>": 51001,
+  "<|12.76|>": 51002,
+  "<|12.78|>": 51003,
+  "<|12.80|>": 51004,
+  "<|12.82|>": 51005,
+  "<|12.84|>": 51006,
+  "<|12.86|>": 51007,
+  "<|12.88|>": 51008,
+  "<|12.90|>": 51009,
+  "<|12.92|>": 51010,
+  "<|12.94|>": 51011,
+  "<|12.96|>": 51012,
+  "<|12.98|>": 51013,
+  "<|13.00|>": 51014,
+  "<|13.02|>": 51015,
+  "<|13.04|>": 51016,
+  "<|13.06|>": 51017,
+  "<|13.08|>": 51018,
+  "<|13.10|>": 51019,
+  "<|13.12|>": 51020,
+  "<|13.14|>": 51021,
+  "<|13.16|>": 51022,
+  "<|13.18|>": 51023,
+  "<|13.20|>": 51024,
+  "<|13.22|>": 51025,
+  "<|13.24|>": 51026,
+  "<|13.26|>": 51027,
+  "<|13.28|>": 51028,
+  "<|13.30|>": 51029,
+  "<|13.32|>": 51030,
+  "<|13.34|>": 51031,
+  "<|13.36|>": 51032,
+  "<|13.38|>": 51033,
+  "<|13.40|>": 51034,
+  "<|13.42|>": 51035,
+  "<|13.44|>": 51036,
+  "<|13.46|>": 51037,
+  "<|13.48|>": 51038,
+  "<|13.50|>": 51039,
+  "<|13.52|>": 51040,
+  "<|13.54|>": 51041,
+  "<|13.56|>": 51042,
+  "<|13.58|>": 51043,
+  "<|13.60|>": 51044,
+  "<|13.62|>": 51045,
+  "<|13.64|>": 51046,
+  "<|13.66|>": 51047,
+  "<|13.68|>": 51048,
+  "<|13.70|>": 51049,
+  "<|13.72|>": 51050,
+  "<|13.74|>": 51051,
+  "<|13.76|>": 51052,
+  "<|13.78|>": 51053,
+  "<|13.80|>": 51054,
+  "<|13.82|>": 51055,
+  "<|13.84|>": 51056,
+  "<|13.86|>": 51057,
+  "<|13.88|>": 51058,
+  "<|13.90|>": 51059,
+  "<|13.92|>": 51060,
+  "<|13.94|>": 51061,
+  "<|13.96|>": 51062,
+  "<|13.98|>": 51063,
+  "<|14.00|>": 51064,
+  "<|14.02|>": 51065,
+  "<|14.04|>": 51066,
+  "<|14.06|>": 51067,
+  "<|14.08|>": 51068,
+  "<|14.10|>": 51069,
+  "<|14.12|>": 51070,
+  "<|14.14|>": 51071,
+  "<|14.16|>": 51072,
+  "<|14.18|>": 51073,
+  "<|14.20|>": 51074,
+  "<|14.22|>": 51075,
+  "<|14.24|>": 51076,
+  "<|14.26|>": 51077,
+  "<|14.28|>": 51078,
+  "<|14.30|>": 51079,
+  "<|14.32|>": 51080,
+  "<|14.34|>": 51081,
+  "<|14.36|>": 51082,
+  "<|14.38|>": 51083,
+  "<|14.40|>": 51084,
+  "<|14.42|>": 51085,
+  "<|14.44|>": 51086,
+  "<|14.46|>": 51087,
+  "<|14.48|>": 51088,
+  "<|14.50|>": 51089,
+  "<|14.52|>": 51090,
+  "<|14.54|>": 51091,
+  "<|14.56|>": 51092,
+  "<|14.58|>": 51093,
+  "<|14.60|>": 51094,
+  "<|14.62|>": 51095,
+  "<|14.64|>": 51096,
+  "<|14.66|>": 51097,
+  "<|14.68|>": 51098,
+  "<|14.70|>": 51099,
+  "<|14.72|>": 51100,
+  "<|14.74|>": 51101,
+  "<|14.76|>": 51102,
+  "<|14.78|>": 51103,
+  "<|14.80|>": 51104,
+  "<|14.82|>": 51105,
+  "<|14.84|>": 51106,
+  "<|14.86|>": 51107,
+  "<|14.88|>": 51108,
+  "<|14.90|>": 51109,
+  "<|14.92|>": 51110,
+  "<|14.94|>": 51111,
+  "<|14.96|>": 51112,
+  "<|14.98|>": 51113,
+  "<|15.00|>": 51114,
+  "<|15.02|>": 51115,
+  "<|15.04|>": 51116,
+  "<|15.06|>": 51117,
+  "<|15.08|>": 51118,
+  "<|15.10|>": 51119,
+  "<|15.12|>": 51120,
+  "<|15.14|>": 51121,
+  "<|15.16|>": 51122,
+  "<|15.18|>": 51123,
+  "<|15.20|>": 51124,
+  "<|15.22|>": 51125,
+  "<|15.24|>": 51126,
+  "<|15.26|>": 51127,
+  "<|15.28|>": 51128,
+  "<|15.30|>": 51129,
+  "<|15.32|>": 51130,
+  "<|15.34|>": 51131,
+  "<|15.36|>": 51132,
+  "<|15.38|>": 51133,
+  "<|15.40|>": 51134,
+  "<|15.42|>": 51135,
+  "<|15.44|>": 51136,
+  "<|15.46|>": 51137,
+  "<|15.48|>": 51138,
+  "<|15.50|>": 51139,
+  "<|15.52|>": 51140,
+  "<|15.54|>": 51141,
+  "<|15.56|>": 51142,
+  "<|15.58|>": 51143,
+  "<|15.60|>": 51144,
+  "<|15.62|>": 51145,
+  "<|15.64|>": 51146,
+  "<|15.66|>": 51147,
+  "<|15.68|>": 51148,
+  "<|15.70|>": 51149,
+  "<|15.72|>": 51150,
+  "<|15.74|>": 51151,
+  "<|15.76|>": 51152,
+  "<|15.78|>": 51153,
+  "<|15.80|>": 51154,
+  "<|15.82|>": 51155,
+  "<|15.84|>": 51156,
+  "<|15.86|>": 51157,
+  "<|15.88|>": 51158,
+  "<|15.90|>": 51159,
+  "<|15.92|>": 51160,
+  "<|15.94|>": 51161,
+  "<|15.96|>": 51162,
+  "<|15.98|>": 51163,
+  "<|16.00|>": 51164,
+  "<|16.02|>": 51165,
+  "<|16.04|>": 51166,
+  "<|16.06|>": 51167,
+  "<|16.08|>": 51168,
+  "<|16.10|>": 51169,
+  "<|16.12|>": 51170,
+  "<|16.14|>": 51171,
+  "<|16.16|>": 51172,
+  "<|16.18|>": 51173,
+  "<|16.20|>": 51174,
+  "<|16.22|>": 51175,
+  "<|16.24|>": 51176,
+  "<|16.26|>": 51177,
+  "<|16.28|>": 51178,
+  "<|16.30|>": 51179,
+  "<|16.32|>": 51180,
+  "<|16.34|>": 51181,
+  "<|16.36|>": 51182,
+  "<|16.38|>": 51183,
+  "<|16.40|>": 51184,
+  "<|16.42|>": 51185,
+  "<|16.44|>": 51186,
+  "<|16.46|>": 51187,
+  "<|16.48|>": 51188,
+  "<|16.50|>": 51189,
+  "<|16.52|>": 51190,
+  "<|16.54|>": 51191,
+  "<|16.56|>": 51192,
+  "<|16.58|>": 51193,
+  "<|16.60|>": 51194,
+  "<|16.62|>": 51195,
+  "<|16.64|>": 51196,
+  "<|16.66|>": 51197,
+  "<|16.68|>": 51198,
+  "<|16.70|>": 51199,
+  "<|16.72|>": 51200,
+  "<|16.74|>": 51201,
+  "<|16.76|>": 51202,
+  "<|16.78|>": 51203,
+  "<|16.80|>": 51204,
+  "<|16.82|>": 51205,
+  "<|16.84|>": 51206,
+  "<|16.86|>": 51207,
+  "<|16.88|>": 51208,
+  "<|16.90|>": 51209,
+  "<|16.92|>": 51210,
+  "<|16.94|>": 51211,
+  "<|16.96|>": 51212,
+  "<|16.98|>": 51213,
+  "<|17.00|>": 51214,
+  "<|17.02|>": 51215,
+  "<|17.04|>": 51216,
+  "<|17.06|>": 51217,
+  "<|17.08|>": 51218,
+  "<|17.10|>": 51219,
+  "<|17.12|>": 51220,
+  "<|17.14|>": 51221,
+  "<|17.16|>": 51222,
+  "<|17.18|>": 51223,
+  "<|17.20|>": 51224,
+  "<|17.22|>": 51225,
+  "<|17.24|>": 51226,
+  "<|17.26|>": 51227,
+  "<|17.28|>": 51228,
+  "<|17.30|>": 51229,
+  "<|17.32|>": 51230,
+  "<|17.34|>": 51231,
+  "<|17.36|>": 51232,
+  "<|17.38|>": 51233,
+  "<|17.40|>": 51234,
+  "<|17.42|>": 51235,
+  "<|17.44|>": 51236,
+  "<|17.46|>": 51237,
+  "<|17.48|>": 51238,
+  "<|17.50|>": 51239,
+  "<|17.52|>": 51240,
+  "<|17.54|>": 51241,
+  "<|17.56|>": 51242,
+  "<|17.58|>": 51243,
+  "<|17.60|>": 51244,
+  "<|17.62|>": 51245,
+  "<|17.64|>": 51246,
+  "<|17.66|>": 51247,
+  "<|17.68|>": 51248,
+  "<|17.70|>": 51249,
+  "<|17.72|>": 51250,
+  "<|17.74|>": 51251,
+  "<|17.76|>": 51252,
+  "<|17.78|>": 51253,
+  "<|17.80|>": 51254,
+  "<|17.82|>": 51255,
+  "<|17.84|>": 51256,
+  "<|17.86|>": 51257,
+  "<|17.88|>": 51258,
+  "<|17.90|>": 51259,
+  "<|17.92|>": 51260,
+  "<|17.94|>": 51261,
+  "<|17.96|>": 51262,
+  "<|17.98|>": 51263,
+  "<|18.00|>": 51264,
+  "<|18.02|>": 51265,
+  "<|18.04|>": 51266,
+  "<|18.06|>": 51267,
+  "<|18.08|>": 51268,
+  "<|18.10|>": 51269,
+  "<|18.12|>": 51270,
+  "<|18.14|>": 51271,
+  "<|18.16|>": 51272,
+  "<|18.18|>": 51273,
+  "<|18.20|>": 51274,
+  "<|18.22|>": 51275,
+  "<|18.24|>": 51276,
+  "<|18.26|>": 51277,
+  "<|18.28|>": 51278,
+  "<|18.30|>": 51279,
+  "<|18.32|>": 51280,
+  "<|18.34|>": 51281,
+  "<|18.36|>": 51282,
+  "<|18.38|>": 51283,
+  "<|18.40|>": 51284,
+  "<|18.42|>": 51285,
+  "<|18.44|>": 51286,
+  "<|18.46|>": 51287,
+  "<|18.48|>": 51288,
+  "<|18.50|>": 51289,
+  "<|18.52|>": 51290,
+  "<|18.54|>": 51291,
+  "<|18.56|>": 51292,
+  "<|18.58|>": 51293,
+  "<|18.60|>": 51294,
+  "<|18.62|>": 51295,
+  "<|18.64|>": 51296,
+  "<|18.66|>": 51297,
+  "<|18.68|>": 51298,
+  "<|18.70|>": 51299,
+  "<|18.72|>": 51300,
+  "<|18.74|>": 51301,
+  "<|18.76|>": 51302,
+  "<|18.78|>": 51303,
+  "<|18.80|>": 51304,
+  "<|18.82|>": 51305,
+  "<|18.84|>": 51306,
+  "<|18.86|>": 51307,
+  "<|18.88|>": 51308,
+  "<|18.90|>": 51309,
+  "<|18.92|>": 51310,
+  "<|18.94|>": 51311,
+  "<|18.96|>": 51312,
+  "<|18.98|>": 51313,
+  "<|19.00|>": 51314,
+  "<|19.02|>": 51315,
+  "<|19.04|>": 51316,
+  "<|19.06|>": 51317,
+  "<|19.08|>": 51318,
+  "<|19.10|>": 51319,
+  "<|19.12|>": 51320,
+  "<|19.14|>": 51321,
+  "<|19.16|>": 51322,
+  "<|19.18|>": 51323,
+  "<|19.20|>": 51324,
+  "<|19.22|>": 51325,
+  "<|19.24|>": 51326,
+  "<|19.26|>": 51327,
+  "<|19.28|>": 51328,
+  "<|19.30|>": 51329,
+  "<|19.32|>": 51330,
+  "<|19.34|>": 51331,
+  "<|19.36|>": 51332,
+  "<|19.38|>": 51333,
+  "<|19.40|>": 51334,
+  "<|19.42|>": 51335,
+  "<|19.44|>": 51336,
+  "<|19.46|>": 51337,
+  "<|19.48|>": 51338,
+  "<|19.50|>": 51339,
+  "<|19.52|>": 51340,
+  "<|19.54|>": 51341,
+  "<|19.56|>": 51342,
+  "<|19.58|>": 51343,
+  "<|19.60|>": 51344,
+  "<|19.62|>": 51345,
+  "<|19.64|>": 51346,
+  "<|19.66|>": 51347,
+  "<|19.68|>": 51348,
+  "<|19.70|>": 51349,
+  "<|19.72|>": 51350,
+  "<|19.74|>": 51351,
+  "<|19.76|>": 51352,
+  "<|19.78|>": 51353,
+  "<|19.80|>": 51354,
+  "<|19.82|>": 51355,
+  "<|19.84|>": 51356,
+  "<|19.86|>": 51357,
+  "<|19.88|>": 51358,
+  "<|19.90|>": 51359,
+  "<|19.92|>": 51360,
+  "<|19.94|>": 51361,
+  "<|19.96|>": 51362,
+  "<|19.98|>": 51363,
+  "<|2.00|>": 50464,
+  "<|2.02|>": 50465,
+  "<|2.04|>": 50466,
+  "<|2.06|>": 50467,
+  "<|2.08|>": 50468,
+  "<|2.10|>": 50469,
+  "<|2.12|>": 50470,
+  "<|2.14|>": 50471,
+  "<|2.16|>": 50472,
+  "<|2.18|>": 50473,
+  "<|2.20|>": 50474,
+  "<|2.22|>": 50475,
+  "<|2.24|>": 50476,
+  "<|2.26|>": 50477,
+  "<|2.28|>": 50478,
+  "<|2.30|>": 50479,
+  "<|2.32|>": 50480,
+  "<|2.34|>": 50481,
+  "<|2.36|>": 50482,
+  "<|2.38|>": 50483,
+  "<|2.40|>": 50484,
+  "<|2.42|>": 50485,
+  "<|2.44|>": 50486,
+  "<|2.46|>": 50487,
+  "<|2.48|>": 50488,
+  "<|2.50|>": 50489,
+  "<|2.52|>": 50490,
+  "<|2.54|>": 50491,
+  "<|2.56|>": 50492,
+  "<|2.58|>": 50493,
+  "<|2.60|>": 50494,
+  "<|2.62|>": 50495,
+  "<|2.64|>": 50496,
+  "<|2.66|>": 50497,
+  "<|2.68|>": 50498,
+  "<|2.70|>": 50499,
+  "<|2.72|>": 50500,
+  "<|2.74|>": 50501,
+  "<|2.76|>": 50502,
+  "<|2.78|>": 50503,
+  "<|2.80|>": 50504,
+  "<|2.82|>": 50505,
+  "<|2.84|>": 50506,
+  "<|2.86|>": 50507,
+  "<|2.88|>": 50508,
+  "<|2.90|>": 50509,
+  "<|2.92|>": 50510,
+  "<|2.94|>": 50511,
+  "<|2.96|>": 50512,
+  "<|2.98|>": 50513,
+  "<|20.00|>": 51364,
+  "<|20.02|>": 51365,
+  "<|20.04|>": 51366,
+  "<|20.06|>": 51367,
+  "<|20.08|>": 51368,
+  "<|20.10|>": 51369,
+  "<|20.12|>": 51370,
+  "<|20.14|>": 51371,
+  "<|20.16|>": 51372,
+  "<|20.18|>": 51373,
+  "<|20.20|>": 51374,
+  "<|20.22|>": 51375,
+  "<|20.24|>": 51376,
+  "<|20.26|>": 51377,
+  "<|20.28|>": 51378,
+  "<|20.30|>": 51379,
+  "<|20.32|>": 51380,
+  "<|20.34|>": 51381,
+  "<|20.36|>": 51382,
+  "<|20.38|>": 51383,
+  "<|20.40|>": 51384,
+  "<|20.42|>": 51385,
+  "<|20.44|>": 51386,
+  "<|20.46|>": 51387,
+  "<|20.48|>": 51388,
+  "<|20.50|>": 51389,
+  "<|20.52|>": 51390,
+  "<|20.54|>": 51391,
+  "<|20.56|>": 51392,
+  "<|20.58|>": 51393,
+  "<|20.60|>": 51394,
+  "<|20.62|>": 51395,
+  "<|20.64|>": 51396,
+  "<|20.66|>": 51397,
+  "<|20.68|>": 51398,
+  "<|20.70|>": 51399,
+  "<|20.72|>": 51400,
+  "<|20.74|>": 51401,
+  "<|20.76|>": 51402,
+  "<|20.78|>": 51403,
+  "<|20.80|>": 51404,
+  "<|20.82|>": 51405,
+  "<|20.84|>": 51406,
+  "<|20.86|>": 51407,
+  "<|20.88|>": 51408,
+  "<|20.90|>": 51409,
+  "<|20.92|>": 51410,
+  "<|20.94|>": 51411,
+  "<|20.96|>": 51412,
+  "<|20.98|>": 51413,
+  "<|21.00|>": 51414,
+  "<|21.02|>": 51415,
+  "<|21.04|>": 51416,
+  "<|21.06|>": 51417,
+  "<|21.08|>": 51418,
+  "<|21.10|>": 51419,
+  "<|21.12|>": 51420,
+  "<|21.14|>": 51421,
+  "<|21.16|>": 51422,
+  "<|21.18|>": 51423,
+  "<|21.20|>": 51424,
+  "<|21.22|>": 51425,
+  "<|21.24|>": 51426,
+  "<|21.26|>": 51427,
+  "<|21.28|>": 51428,
+  "<|21.30|>": 51429,
+  "<|21.32|>": 51430,
+  "<|21.34|>": 51431,
+  "<|21.36|>": 51432,
+  "<|21.38|>": 51433,
+  "<|21.40|>": 51434,
+  "<|21.42|>": 51435,
+  "<|21.44|>": 51436,
+  "<|21.46|>": 51437,
+  "<|21.48|>": 51438,
+  "<|21.50|>": 51439,
+  "<|21.52|>": 51440,
+  "<|21.54|>": 51441,
+  "<|21.56|>": 51442,
+  "<|21.58|>": 51443,
+  "<|21.60|>": 51444,
+  "<|21.62|>": 51445,
+  "<|21.64|>": 51446,
+  "<|21.66|>": 51447,
+  "<|21.68|>": 51448,
+  "<|21.70|>": 51449,
+  "<|21.72|>": 51450,
+  "<|21.74|>": 51451,
+  "<|21.76|>": 51452,
+  "<|21.78|>": 51453,
+  "<|21.80|>": 51454,
+  "<|21.82|>": 51455,
+  "<|21.84|>": 51456,
+  "<|21.86|>": 51457,
+  "<|21.88|>": 51458,
+  "<|21.90|>": 51459,
+  "<|21.92|>": 51460,
+  "<|21.94|>": 51461,
+  "<|21.96|>": 51462,
+  "<|21.98|>": 51463,
+  "<|22.00|>": 51464,
+  "<|22.02|>": 51465,
+  "<|22.04|>": 51466,
+  "<|22.06|>": 51467,
+  "<|22.08|>": 51468,
+  "<|22.10|>": 51469,
+  "<|22.12|>": 51470,
+  "<|22.14|>": 51471,
+  "<|22.16|>": 51472,
+  "<|22.18|>": 51473,
+  "<|22.20|>": 51474,
+  "<|22.22|>": 51475,
+  "<|22.24|>": 51476,
+  "<|22.26|>": 51477,
+  "<|22.28|>": 51478,
+  "<|22.30|>": 51479,
+  "<|22.32|>": 51480,
+  "<|22.34|>": 51481,
+  "<|22.36|>": 51482,
+  "<|22.38|>": 51483,
+  "<|22.40|>": 51484,
+  "<|22.42|>": 51485,
+  "<|22.44|>": 51486,
+  "<|22.46|>": 51487,
+  "<|22.48|>": 51488,
+  "<|22.50|>": 51489,
+  "<|22.52|>": 51490,
+  "<|22.54|>": 51491,
+  "<|22.56|>": 51492,
+  "<|22.58|>": 51493,
+  "<|22.60|>": 51494,
+  "<|22.62|>": 51495,
+  "<|22.64|>": 51496,
+  "<|22.66|>": 51497,
+  "<|22.68|>": 51498,
+  "<|22.70|>": 51499,
+  "<|22.72|>": 51500,
+  "<|22.74|>": 51501,
+  "<|22.76|>": 51502,
+  "<|22.78|>": 51503,
+  "<|22.80|>": 51504,
+  "<|22.82|>": 51505,
+  "<|22.84|>": 51506,
+  "<|22.86|>": 51507,
+  "<|22.88|>": 51508,
+  "<|22.90|>": 51509,
+  "<|22.92|>": 51510,
+  "<|22.94|>": 51511,
+  "<|22.96|>": 51512,
+  "<|22.98|>": 51513,
+  "<|23.00|>": 51514,
+  "<|23.02|>": 51515,
+  "<|23.04|>": 51516,
+  "<|23.06|>": 51517,
+  "<|23.08|>": 51518,
+  "<|23.10|>": 51519,
+  "<|23.12|>": 51520,
+  "<|23.14|>": 51521,
+  "<|23.16|>": 51522,
+  "<|23.18|>": 51523,
+  "<|23.20|>": 51524,
+  "<|23.22|>": 51525,
+  "<|23.24|>": 51526,
+  "<|23.26|>": 51527,
+  "<|23.28|>": 51528,
+  "<|23.30|>": 51529,
+  "<|23.32|>": 51530,
+  "<|23.34|>": 51531,
+  "<|23.36|>": 51532,
+  "<|23.38|>": 51533,
+  "<|23.40|>": 51534,
+  "<|23.42|>": 51535,
+  "<|23.44|>": 51536,
+  "<|23.46|>": 51537,
+  "<|23.48|>": 51538,
+  "<|23.50|>": 51539,
+  "<|23.52|>": 51540,
+  "<|23.54|>": 51541,
+  "<|23.56|>": 51542,
+  "<|23.58|>": 51543,
+  "<|23.60|>": 51544,
+  "<|23.62|>": 51545,
+  "<|23.64|>": 51546,
+  "<|23.66|>": 51547,
+  "<|23.68|>": 51548,
+  "<|23.70|>": 51549,
+  "<|23.72|>": 51550,
+  "<|23.74|>": 51551,
+  "<|23.76|>": 51552,
+  "<|23.78|>": 51553,
+  "<|23.80|>": 51554,
+  "<|23.82|>": 51555,
+  "<|23.84|>": 51556,
+  "<|23.86|>": 51557,
+  "<|23.88|>": 51558,
+  "<|23.90|>": 51559,
+  "<|23.92|>": 51560,
+  "<|23.94|>": 51561,
+  "<|23.96|>": 51562,
+  "<|23.98|>": 51563,
+  "<|24.00|>": 51564,
+  "<|24.02|>": 51565,
+  "<|24.04|>": 51566,
+  "<|24.06|>": 51567,
+  "<|24.08|>": 51568,
+  "<|24.10|>": 51569,
+  "<|24.12|>": 51570,
+  "<|24.14|>": 51571,
+  "<|24.16|>": 51572,
+  "<|24.18|>": 51573,
+  "<|24.20|>": 51574,
+  "<|24.22|>": 51575,
+  "<|24.24|>": 51576,
+  "<|24.26|>": 51577,
+  "<|24.28|>": 51578,
+  "<|24.30|>": 51579,
+  "<|24.32|>": 51580,
+  "<|24.34|>": 51581,
+  "<|24.36|>": 51582,
+  "<|24.38|>": 51583,
+  "<|24.40|>": 51584,
+  "<|24.42|>": 51585,
+  "<|24.44|>": 51586,
+  "<|24.46|>": 51587,
+  "<|24.48|>": 51588,
+  "<|24.50|>": 51589,
+  "<|24.52|>": 51590,
+  "<|24.54|>": 51591,
+  "<|24.56|>": 51592,
+  "<|24.58|>": 51593,
+  "<|24.60|>": 51594,
+  "<|24.62|>": 51595,
+  "<|24.64|>": 51596,
+  "<|24.66|>": 51597,
+  "<|24.68|>": 51598,
+  "<|24.70|>": 51599,
+  "<|24.72|>": 51600,
+  "<|24.74|>": 51601,
+  "<|24.76|>": 51602,
+  "<|24.78|>": 51603,
+  "<|24.80|>": 51604,
+  "<|24.82|>": 51605,
+  "<|24.84|>": 51606,
+  "<|24.86|>": 51607,
+  "<|24.88|>": 51608,
+  "<|24.90|>": 51609,
+  "<|24.92|>": 51610,
+  "<|24.94|>": 51611,
+  "<|24.96|>": 51612,
+  "<|24.98|>": 51613,
+  "<|25.00|>": 51614,
+  "<|25.02|>": 51615,
+  "<|25.04|>": 51616,
+  "<|25.06|>": 51617,
+  "<|25.08|>": 51618,
+  "<|25.10|>": 51619,
+  "<|25.12|>": 51620,
+  "<|25.14|>": 51621,
+  "<|25.16|>": 51622,
+  "<|25.18|>": 51623,
+  "<|25.20|>": 51624,
+  "<|25.22|>": 51625,
+  "<|25.24|>": 51626,
+  "<|25.26|>": 51627,
+  "<|25.28|>": 51628,
+  "<|25.30|>": 51629,
+  "<|25.32|>": 51630,
+  "<|25.34|>": 51631,
+  "<|25.36|>": 51632,
+  "<|25.38|>": 51633,
+  "<|25.40|>": 51634,
+  "<|25.42|>": 51635,
+  "<|25.44|>": 51636,
+  "<|25.46|>": 51637,
+  "<|25.48|>": 51638,
+  "<|25.50|>": 51639,
+  "<|25.52|>": 51640,
+  "<|25.54|>": 51641,
+  "<|25.56|>": 51642,
+  "<|25.58|>": 51643,
+  "<|25.60|>": 51644,
+  "<|25.62|>": 51645,
+  "<|25.64|>": 51646,
+  "<|25.66|>": 51647,
+  "<|25.68|>": 51648,
+  "<|25.70|>": 51649,
+  "<|25.72|>": 51650,
+  "<|25.74|>": 51651,
+  "<|25.76|>": 51652,
+  "<|25.78|>": 51653,
+  "<|25.80|>": 51654,
+  "<|25.82|>": 51655,
+  "<|25.84|>": 51656,
+  "<|25.86|>": 51657,
+  "<|25.88|>": 51658,
+  "<|25.90|>": 51659,
+  "<|25.92|>": 51660,
+  "<|25.94|>": 51661,
+  "<|25.96|>": 51662,
+  "<|25.98|>": 51663,
+  "<|26.00|>": 51664,
+  "<|26.02|>": 51665,
+  "<|26.04|>": 51666,
+  "<|26.06|>": 51667,
+  "<|26.08|>": 51668,
+  "<|26.10|>": 51669,
+  "<|26.12|>": 51670,
+  "<|26.14|>": 51671,
+  "<|26.16|>": 51672,
+  "<|26.18|>": 51673,
+  "<|26.20|>": 51674,
+  "<|26.22|>": 51675,
+  "<|26.24|>": 51676,
+  "<|26.26|>": 51677,
+  "<|26.28|>": 51678,
+  "<|26.30|>": 51679,
+  "<|26.32|>": 51680,
+  "<|26.34|>": 51681,
+  "<|26.36|>": 51682,
+  "<|26.38|>": 51683,
+  "<|26.40|>": 51684,
+  "<|26.42|>": 51685,
+  "<|26.44|>": 51686,
+  "<|26.46|>": 51687,
+  "<|26.48|>": 51688,
+  "<|26.50|>": 51689,
+  "<|26.52|>": 51690,
+  "<|26.54|>": 51691,
+  "<|26.56|>": 51692,
+  "<|26.58|>": 51693,
+  "<|26.60|>": 51694,
+  "<|26.62|>": 51695,
+  "<|26.64|>": 51696,
+  "<|26.66|>": 51697,
+  "<|26.68|>": 51698,
+  "<|26.70|>": 51699,
+  "<|26.72|>": 51700,
+  "<|26.74|>": 51701,
+  "<|26.76|>": 51702,
+  "<|26.78|>": 51703,
+  "<|26.80|>": 51704,
+  "<|26.82|>": 51705,
+  "<|26.84|>": 51706,
+  "<|26.86|>": 51707,
+  "<|26.88|>": 51708,
+  "<|26.90|>": 51709,
+  "<|26.92|>": 51710,
+  "<|26.94|>": 51711,
+  "<|26.96|>": 51712,
+  "<|26.98|>": 51713,
+  "<|27.00|>": 51714,
+  "<|27.02|>": 51715,
+  "<|27.04|>": 51716,
+  "<|27.06|>": 51717,
+  "<|27.08|>": 51718,
+  "<|27.10|>": 51719,
+  "<|27.12|>": 51720,
+  "<|27.14|>": 51721,
+  "<|27.16|>": 51722,
+  "<|27.18|>": 51723,
+  "<|27.20|>": 51724,
+  "<|27.22|>": 51725,
+  "<|27.24|>": 51726,
+  "<|27.26|>": 51727,
+  "<|27.28|>": 51728,
+  "<|27.30|>": 51729,
+  "<|27.32|>": 51730,
+  "<|27.34|>": 51731,
+  "<|27.36|>": 51732,
+  "<|27.38|>": 51733,
+  "<|27.40|>": 51734,
+  "<|27.42|>": 51735,
+  "<|27.44|>": 51736,
+  "<|27.46|>": 51737,
+  "<|27.48|>": 51738,
+  "<|27.50|>": 51739,
+  "<|27.52|>": 51740,
+  "<|27.54|>": 51741,
+  "<|27.56|>": 51742,
+  "<|27.58|>": 51743,
+  "<|27.60|>": 51744,
+  "<|27.62|>": 51745,
+  "<|27.64|>": 51746,
+  "<|27.66|>": 51747,
+  "<|27.68|>": 51748,
+  "<|27.70|>": 51749,
+  "<|27.72|>": 51750,
+  "<|27.74|>": 51751,
+  "<|27.76|>": 51752,
+  "<|27.78|>": 51753,
+  "<|27.80|>": 51754,
+  "<|27.82|>": 51755,
+  "<|27.84|>": 51756,
+  "<|27.86|>": 51757,
+  "<|27.88|>": 51758,
+  "<|27.90|>": 51759,
+  "<|27.92|>": 51760,
+  "<|27.94|>": 51761,
+  "<|27.96|>": 51762,
+  "<|27.98|>": 51763,
+  "<|28.00|>": 51764,
+  "<|28.02|>": 51765,
+  "<|28.04|>": 51766,
+  "<|28.06|>": 51767,
+  "<|28.08|>": 51768,
+  "<|28.10|>": 51769,
+  "<|28.12|>": 51770,
+  "<|28.14|>": 51771,
+  "<|28.16|>": 51772,
+  "<|28.18|>": 51773,
+  "<|28.20|>": 51774,
+  "<|28.22|>": 51775,
+  "<|28.24|>": 51776,
+  "<|28.26|>": 51777,
+  "<|28.28|>": 51778,
+  "<|28.30|>": 51779,
+  "<|28.32|>": 51780,
+  "<|28.34|>": 51781,
+  "<|28.36|>": 51782,
+  "<|28.38|>": 51783,
+  "<|28.40|>": 51784,
+  "<|28.42|>": 51785,
+  "<|28.44|>": 51786,
+  "<|28.46|>": 51787,
+  "<|28.48|>": 51788,
+  "<|28.50|>": 51789,
+  "<|28.52|>": 51790,
+  "<|28.54|>": 51791,
+  "<|28.56|>": 51792,
+  "<|28.58|>": 51793,
+  "<|28.60|>": 51794,
+  "<|28.62|>": 51795,
+  "<|28.64|>": 51796,
+  "<|28.66|>": 51797,
+  "<|28.68|>": 51798,
+  "<|28.70|>": 51799,
+  "<|28.72|>": 51800,
+  "<|28.74|>": 51801,
+  "<|28.76|>": 51802,
+  "<|28.78|>": 51803,
+  "<|28.80|>": 51804,
+  "<|28.82|>": 51805,
+  "<|28.84|>": 51806,
+  "<|28.86|>": 51807,
+  "<|28.88|>": 51808,
+  "<|28.90|>": 51809,
+  "<|28.92|>": 51810,
+  "<|28.94|>": 51811,
+  "<|28.96|>": 51812,
+  "<|28.98|>": 51813,
+  "<|29.00|>": 51814,
+  "<|29.02|>": 51815,
+  "<|29.04|>": 51816,
+  "<|29.06|>": 51817,
+  "<|29.08|>": 51818,
+  "<|29.10|>": 51819,
+  "<|29.12|>": 51820,
+  "<|29.14|>": 51821,
+  "<|29.16|>": 51822,
+  "<|29.18|>": 51823,
+  "<|29.20|>": 51824,
+  "<|29.22|>": 51825,
+  "<|29.24|>": 51826,
+  "<|29.26|>": 51827,
+  "<|29.28|>": 51828,
+  "<|29.30|>": 51829,
+  "<|29.32|>": 51830,
+  "<|29.34|>": 51831,
+  "<|29.36|>": 51832,
+  "<|29.38|>": 51833,
+  "<|29.40|>": 51834,
+  "<|29.42|>": 51835,
+  "<|29.44|>": 51836,
+  "<|29.46|>": 51837,
+  "<|29.48|>": 51838,
+  "<|29.50|>": 51839,
+  "<|29.52|>": 51840,
+  "<|29.54|>": 51841,
+  "<|29.56|>": 51842,
+  "<|29.58|>": 51843,
+  "<|29.60|>": 51844,
+  "<|29.62|>": 51845,
+  "<|29.64|>": 51846,
+  "<|29.66|>": 51847,
+  "<|29.68|>": 51848,
+  "<|29.70|>": 51849,
+  "<|29.72|>": 51850,
+  "<|29.74|>": 51851,
+  "<|29.76|>": 51852,
+  "<|29.78|>": 51853,
+  "<|29.80|>": 51854,
+  "<|29.82|>": 51855,
+  "<|29.84|>": 51856,
+  "<|29.86|>": 51857,
+  "<|29.88|>": 51858,
+  "<|29.90|>": 51859,
+  "<|29.92|>": 51860,
+  "<|29.94|>": 51861,
+  "<|29.96|>": 51862,
+  "<|29.98|>": 51863,
+  "<|3.00|>": 50514,
+  "<|3.02|>": 50515,
+  "<|3.04|>": 50516,
+  "<|3.06|>": 50517,
+  "<|3.08|>": 50518,
+  "<|3.10|>": 50519,
+  "<|3.12|>": 50520,
+  "<|3.14|>": 50521,
+  "<|3.16|>": 50522,
+  "<|3.18|>": 50523,
+  "<|3.20|>": 50524,
+  "<|3.22|>": 50525,
+  "<|3.24|>": 50526,
+  "<|3.26|>": 50527,
+  "<|3.28|>": 50528,
+  "<|3.30|>": 50529,
+  "<|3.32|>": 50530,
+  "<|3.34|>": 50531,
+  "<|3.36|>": 50532,
+  "<|3.38|>": 50533,
+  "<|3.40|>": 50534,
+  "<|3.42|>": 50535,
+  "<|3.44|>": 50536,
+  "<|3.46|>": 50537,
+  "<|3.48|>": 50538,
+  "<|3.50|>": 50539,
+  "<|3.52|>": 50540,
+  "<|3.54|>": 50541,
+  "<|3.56|>": 50542,
+  "<|3.58|>": 50543,
+  "<|3.60|>": 50544,
+  "<|3.62|>": 50545,
+  "<|3.64|>": 50546,
+  "<|3.66|>": 50547,
+  "<|3.68|>": 50548,
+  "<|3.70|>": 50549,
+  "<|3.72|>": 50550,
+  "<|3.74|>": 50551,
+  "<|3.76|>": 50552,
+  "<|3.78|>": 50553,
+  "<|3.80|>": 50554,
+  "<|3.82|>": 50555,
+  "<|3.84|>": 50556,
+  "<|3.86|>": 50557,
+  "<|3.88|>": 50558,
+  "<|3.90|>": 50559,
+  "<|3.92|>": 50560,
+  "<|3.94|>": 50561,
+  "<|3.96|>": 50562,
+  "<|3.98|>": 50563,
+  "<|30.00|>": 51864,
+  "<|4.00|>": 50564,
+  "<|4.02|>": 50565,
+  "<|4.04|>": 50566,
+  "<|4.06|>": 50567,
+  "<|4.08|>": 50568,
+  "<|4.10|>": 50569,
+  "<|4.12|>": 50570,
+  "<|4.14|>": 50571,
+  "<|4.16|>": 50572,
+  "<|4.18|>": 50573,
+  "<|4.20|>": 50574,
+  "<|4.22|>": 50575,
+  "<|4.24|>": 50576,
+  "<|4.26|>": 50577,
+  "<|4.28|>": 50578,
+  "<|4.30|>": 50579,
+  "<|4.32|>": 50580,
+  "<|4.34|>": 50581,
+  "<|4.36|>": 50582,
+  "<|4.38|>": 50583,
+  "<|4.40|>": 50584,
+  "<|4.42|>": 50585,
+  "<|4.44|>": 50586,
+  "<|4.46|>": 50587,
+  "<|4.48|>": 50588,
+  "<|4.50|>": 50589,
+  "<|4.52|>": 50590,
+  "<|4.54|>": 50591,
+  "<|4.56|>": 50592,
+  "<|4.58|>": 50593,
+  "<|4.60|>": 50594,
+  "<|4.62|>": 50595,
+  "<|4.64|>": 50596,
+  "<|4.66|>": 50597,
+  "<|4.68|>": 50598,
+  "<|4.70|>": 50599,
+  "<|4.72|>": 50600,
+  "<|4.74|>": 50601,
+  "<|4.76|>": 50602,
+  "<|4.78|>": 50603,
+  "<|4.80|>": 50604,
+  "<|4.82|>": 50605,
+  "<|4.84|>": 50606,
+  "<|4.86|>": 50607,
+  "<|4.88|>": 50608,
+  "<|4.90|>": 50609,
+  "<|4.92|>": 50610,
+  "<|4.94|>": 50611,
+  "<|4.96|>": 50612,
+  "<|4.98|>": 50613,
+  "<|5.00|>": 50614,
+  "<|5.02|>": 50615,
+  "<|5.04|>": 50616,
+  "<|5.06|>": 50617,
+  "<|5.08|>": 50618,
+  "<|5.10|>": 50619,
+  "<|5.12|>": 50620,
+  "<|5.14|>": 50621,
+  "<|5.16|>": 50622,
+  "<|5.18|>": 50623,
+  "<|5.20|>": 50624,
+  "<|5.22|>": 50625,
+  "<|5.24|>": 50626,
+  "<|5.26|>": 50627,
+  "<|5.28|>": 50628,
+  "<|5.30|>": 50629,
+  "<|5.32|>": 50630,
+  "<|5.34|>": 50631,
+  "<|5.36|>": 50632,
+  "<|5.38|>": 50633,
+  "<|5.40|>": 50634,
+  "<|5.42|>": 50635,
+  "<|5.44|>": 50636,
+  "<|5.46|>": 50637,
+  "<|5.48|>": 50638,
+  "<|5.50|>": 50639,
+  "<|5.52|>": 50640,
+  "<|5.54|>": 50641,
+  "<|5.56|>": 50642,
+  "<|5.58|>": 50643,
+  "<|5.60|>": 50644,
+  "<|5.62|>": 50645,
+  "<|5.64|>": 50646,
+  "<|5.66|>": 50647,
+  "<|5.68|>": 50648,
+  "<|5.70|>": 50649,
+  "<|5.72|>": 50650,
+  "<|5.74|>": 50651,
+  "<|5.76|>": 50652,
+  "<|5.78|>": 50653,
+  "<|5.80|>": 50654,
+  "<|5.82|>": 50655,
+  "<|5.84|>": 50656,
+  "<|5.86|>": 50657,
+  "<|5.88|>": 50658,
+  "<|5.90|>": 50659,
+  "<|5.92|>": 50660,
+  "<|5.94|>": 50661,
+  "<|5.96|>": 50662,
+  "<|5.98|>": 50663,
+  "<|6.00|>": 50664,
+  "<|6.02|>": 50665,
+  "<|6.04|>": 50666,
+  "<|6.06|>": 50667,
+  "<|6.08|>": 50668,
+  "<|6.10|>": 50669,
+  "<|6.12|>": 50670,
+  "<|6.14|>": 50671,
+  "<|6.16|>": 50672,
+  "<|6.18|>": 50673,
+  "<|6.20|>": 50674,
+  "<|6.22|>": 50675,
+  "<|6.24|>": 50676,
+  "<|6.26|>": 50677,
+  "<|6.28|>": 50678,
+  "<|6.30|>": 50679,
+  "<|6.32|>": 50680,
+  "<|6.34|>": 50681,
+  "<|6.36|>": 50682,
+  "<|6.38|>": 50683,
+  "<|6.40|>": 50684,
+  "<|6.42|>": 50685,
+  "<|6.44|>": 50686,
+  "<|6.46|>": 50687,
+  "<|6.48|>": 50688,
+  "<|6.50|>": 50689,
+  "<|6.52|>": 50690,
+  "<|6.54|>": 50691,
+  "<|6.56|>": 50692,
+  "<|6.58|>": 50693,
+  "<|6.60|>": 50694,
+  "<|6.62|>": 50695,
+  "<|6.64|>": 50696,
+  "<|6.66|>": 50697,
+  "<|6.68|>": 50698,
+  "<|6.70|>": 50699,
+  "<|6.72|>": 50700,
+  "<|6.74|>": 50701,
+  "<|6.76|>": 50702,
+  "<|6.78|>": 50703,
+  "<|6.80|>": 50704,
+  "<|6.82|>": 50705,
+  "<|6.84|>": 50706,
+  "<|6.86|>": 50707,
+  "<|6.88|>": 50708,
+  "<|6.90|>": 50709,
+  "<|6.92|>": 50710,
+  "<|6.94|>": 50711,
+  "<|6.96|>": 50712,
+  "<|6.98|>": 50713,
+  "<|7.00|>": 50714,
+  "<|7.02|>": 50715,
+  "<|7.04|>": 50716,
+  "<|7.06|>": 50717,
+  "<|7.08|>": 50718,
+  "<|7.10|>": 50719,
+  "<|7.12|>": 50720,
+  "<|7.14|>": 50721,
+  "<|7.16|>": 50722,
+  "<|7.18|>": 50723,
+  "<|7.20|>": 50724,
+  "<|7.22|>": 50725,
+  "<|7.24|>": 50726,
+  "<|7.26|>": 50727,
+  "<|7.28|>": 50728,
+  "<|7.30|>": 50729,
+  "<|7.32|>": 50730,
+  "<|7.34|>": 50731,
+  "<|7.36|>": 50732,
+  "<|7.38|>": 50733,
+  "<|7.40|>": 50734,
+  "<|7.42|>": 50735,
+  "<|7.44|>": 50736,
+  "<|7.46|>": 50737,
+  "<|7.48|>": 50738,
+  "<|7.50|>": 50739,
+  "<|7.52|>": 50740,
+  "<|7.54|>": 50741,
+  "<|7.56|>": 50742,
+  "<|7.58|>": 50743,
+  "<|7.60|>": 50744,
+  "<|7.62|>": 50745,
+  "<|7.64|>": 50746,
+  "<|7.66|>": 50747,
+  "<|7.68|>": 50748,
+  "<|7.70|>": 50749,
+  "<|7.72|>": 50750,
+  "<|7.74|>": 50751,
+  "<|7.76|>": 50752,
+  "<|7.78|>": 50753,
+  "<|7.80|>": 50754,
+  "<|7.82|>": 50755,
+  "<|7.84|>": 50756,
+  "<|7.86|>": 50757,
+  "<|7.88|>": 50758,
+  "<|7.90|>": 50759,
+  "<|7.92|>": 50760,
+  "<|7.94|>": 50761,
+  "<|7.96|>": 50762,
+  "<|7.98|>": 50763,
+  "<|8.00|>": 50764,
+  "<|8.02|>": 50765,
+  "<|8.04|>": 50766,
+  "<|8.06|>": 50767,
+  "<|8.08|>": 50768,
+  "<|8.10|>": 50769,
+  "<|8.12|>": 50770,
+  "<|8.14|>": 50771,
+  "<|8.16|>": 50772,
+  "<|8.18|>": 50773,
+  "<|8.20|>": 50774,
+  "<|8.22|>": 50775,
+  "<|8.24|>": 50776,
+  "<|8.26|>": 50777,
+  "<|8.28|>": 50778,
+  "<|8.30|>": 50779,
+  "<|8.32|>": 50780,
+  "<|8.34|>": 50781,
+  "<|8.36|>": 50782,
+  "<|8.38|>": 50783,
+  "<|8.40|>": 50784,
+  "<|8.42|>": 50785,
+  "<|8.44|>": 50786,
+  "<|8.46|>": 50787,
+  "<|8.48|>": 50788,
+  "<|8.50|>": 50789,
+  "<|8.52|>": 50790,
+  "<|8.54|>": 50791,
+  "<|8.56|>": 50792,
+  "<|8.58|>": 50793,
+  "<|8.60|>": 50794,
+  "<|8.62|>": 50795,
+  "<|8.64|>": 50796,
+  "<|8.66|>": 50797,
+  "<|8.68|>": 50798,
+  "<|8.70|>": 50799,
+  "<|8.72|>": 50800,
+  "<|8.74|>": 50801,
+  "<|8.76|>": 50802,
+  "<|8.78|>": 50803,
+  "<|8.80|>": 50804,
+  "<|8.82|>": 50805,
+  "<|8.84|>": 50806,
+  "<|8.86|>": 50807,
+  "<|8.88|>": 50808,
+  "<|8.90|>": 50809,
+  "<|8.92|>": 50810,
+  "<|8.94|>": 50811,
+  "<|8.96|>": 50812,
+  "<|8.98|>": 50813,
+  "<|9.00|>": 50814,
+  "<|9.02|>": 50815,
+  "<|9.04|>": 50816,
+  "<|9.06|>": 50817,
+  "<|9.08|>": 50818,
+  "<|9.10|>": 50819,
+  "<|9.12|>": 50820,
+  "<|9.14|>": 50821,
+  "<|9.16|>": 50822,
+  "<|9.18|>": 50823,
+  "<|9.20|>": 50824,
+  "<|9.22|>": 50825,
+  "<|9.24|>": 50826,
+  "<|9.26|>": 50827,
+  "<|9.28|>": 50828,
+  "<|9.30|>": 50829,
+  "<|9.32|>": 50830,
+  "<|9.34|>": 50831,
+  "<|9.36|>": 50832,
+  "<|9.38|>": 50833,
+  "<|9.40|>": 50834,
+  "<|9.42|>": 50835,
+  "<|9.44|>": 50836,
+  "<|9.46|>": 50837,
+  "<|9.48|>": 50838,
+  "<|9.50|>": 50839,
+  "<|9.52|>": 50840,
+  "<|9.54|>": 50841,
+  "<|9.56|>": 50842,
+  "<|9.58|>": 50843,
+  "<|9.60|>": 50844,
+  "<|9.62|>": 50845,
+  "<|9.64|>": 50846,
+  "<|9.66|>": 50847,
+  "<|9.68|>": 50848,
+  "<|9.70|>": 50849,
+  "<|9.72|>": 50850,
+  "<|9.74|>": 50851,
+  "<|9.76|>": 50852,
+  "<|9.78|>": 50853,
+  "<|9.80|>": 50854,
+  "<|9.82|>": 50855,
+  "<|9.84|>": 50856,
+  "<|9.86|>": 50857,
+  "<|9.88|>": 50858,
+  "<|9.90|>": 50859,
+  "<|9.92|>": 50860,
+  "<|9.94|>": 50861,
+  "<|9.96|>": 50862,
+  "<|9.98|>": 50863,
+  "<|af|>": 50327,
+  "<|am|>": 50334,
+  "<|ar|>": 50272,
+  "<|as|>": 50350,
+  "<|az|>": 50304,
+  "<|ba|>": 50355,
+  "<|be|>": 50330,
+  "<|bg|>": 50292,
+  "<|bn|>": 50302,
+  "<|bo|>": 50347,
+  "<|br|>": 50309,
+  "<|bs|>": 50315,
+  "<|ca|>": 50270,
+  "<|cs|>": 50283,
+  "<|cy|>": 50297,
+  "<|da|>": 50285,
+  "<|de|>": 50261,
+  "<|el|>": 50281,
+  "<|en|>": 50259,
+  "<|es|>": 50262,
+  "<|et|>": 50307,
+  "<|eu|>": 50310,
+  "<|fa|>": 50300,
+  "<|fi|>": 50277,
+  "<|fo|>": 50338,
+  "<|fr|>": 50265,
+  "<|gl|>": 50319,
+  "<|gu|>": 50333,
+  "<|haw|>": 50352,
+  "<|ha|>": 50354,
+  "<|he|>": 50279,
+  "<|hi|>": 50276,
+  "<|hr|>": 50291,
+  "<|ht|>": 50339,
+  "<|hu|>": 50286,
+  "<|hy|>": 50312,
+  "<|id|>": 50275,
+  "<|is|>": 50311,
+  "<|it|>": 50274,
+  "<|ja|>": 50266,
+  "<|jw|>": 50356,
+  "<|ka|>": 50329,
+  "<|kk|>": 50316,
+  "<|km|>": 50323,
+  "<|kn|>": 50306,
+  "<|ko|>": 50264,
+  "<|la|>": 50294,
+  "<|lb|>": 50345,
+  "<|ln|>": 50353,
+  "<|lo|>": 50336,
+  "<|lt|>": 50293,
+  "<|lv|>": 50301,
+  "<|mg|>": 50349,
+  "<|mi|>": 50295,
+  "<|mk|>": 50308,
+  "<|ml|>": 50296,
+  "<|mn|>": 50314,
+  "<|mr|>": 50320,
+  "<|ms|>": 50282,
+  "<|mt|>": 50343,
+  "<|my|>": 50346,
+  "<|ne|>": 50313,
+  "<|nl|>": 50271,
+  "<|nn|>": 50342,
+  "<|nocaptions|>": 50362,
+  "<|notimestamps|>": 50363,
+  "<|no|>": 50288,
+  "<|oc|>": 50328,
+  "<|pa|>": 50321,
+  "<|pl|>": 50269,
+  "<|ps|>": 50340,
+  "<|pt|>": 50267,
+  "<|ro|>": 50284,
+  "<|ru|>": 50263,
+  "<|sa|>": 50344,
+  "<|sd|>": 50332,
+  "<|si|>": 50322,
+  "<|sk|>": 50298,
+  "<|sl|>": 50305,
+  "<|sn|>": 50324,
+  "<|so|>": 50326,
+  "<|sq|>": 50317,
+  "<|sr|>": 50303,
+  "<|startoflm|>": 50360,
+  "<|startofprev|>": 50361,
+  "<|startoftranscript|>": 50258,
+  "<|su|>": 50357,
+  "<|sv|>": 50273,
+  "<|sw|>": 50318,
+  "<|ta|>": 50287,
+  "<|te|>": 50299,
+  "<|tg|>": 50331,
+  "<|th|>": 50289,
+  "<|tk|>": 50341,
+  "<|tl|>": 50348,
+  "<|transcribe|>": 50359,
+  "<|translate|>": 50358,
+  "<|tr|>": 50268,
+  "<|tt|>": 50351,
+  "<|uk|>": 50280,
+  "<|ur|>": 50290,
+  "<|uz|>": 50337,
+  "<|vi|>": 50278,
+  "<|yi|>": 50335,
+  "<|yo|>": 50325,
+  "<|zh|>": 50260
+}

computer-vision-study-group/README.md

@@ -0,0 +1,15 @@
+# Computer Vision Study Group
+
+This is a collection of all past sessions that have been held as part of the Hugging Face Computer Vision Study Group.
+
+| |Session Name | Session Link  |
+|--- |--- | --- |
+|❓|How Do Vision Transformers Work? | [Session Sheet](Sessions/HowDoVisionTransformersWork.md) |
+|🔅|Polarized Self-Attention | [Session Sheet](Sessions/PolarizedSelfAttention.md)|
+|🍄|Swin Transformer | [Session Sheet](Sessions/SwinTransformer.md)|
+|🔮|Introduction to Neural Radiance Fields | [Session Sheet](Sessions/NeuralRadianceFields.md)|
+|🌐|Hugging Face Vision Ecosystem Overview (June 2022) | [Session Sheet](Sessions/HFVisionEcosystem.md)|
+|🪂|Masked Autoencoders Are Scalable Vision Learners | [Session Sheet](Sessions/MaskedAutoEncoders.md)|
+|🦊|Fiber: Coarse-to-Fine Vision-Language Pre-Training | [Session Sheet](Sessions/Fiber.md)|
+|⚔️ |FlexiViT: One Model for All Patch Sizes| [Session Sheet](Sessions/FlexiViT.md)|
+|🤖|BLIP-2: Bootstrapping Language-Image Pre-training| [Session Sheet](Sessions/Blip2.md)|

computer-vision-study-group/Sessions/Blip2.md

@@ -0,0 +1,25 @@
+# BLIP-2: Bootstrapping Language-Image Pre-training with Frozen Image Encoders and Large Language Models
+Session by [johko](https://github.com/johko)
+
+
+## Recording 📺
+[YouTube](https://www.youtube.com/watch?v=k0DAtZCCl1w&pp=ygUdaHVnZ2luZyBmYWNlIHN0dWR5IGdyb3VwIHN3aW4%3D)
+
+
+## Session Slides 🖥️
+[Google Drive](https://docs.google.com/presentation/d/1Y_8Qu0CMlt7jvCd8Jw0c_ILh8LHB0XgnlrvXObe5FYs/edit?usp=sharing)
+
+
+## Original Paper 📄
+[Hugging Face](https://huggingface.co/papers/2301.12597) /
+[arxiv](https://arxiv.org/abs/2301.12597)
+
+
+## GitHub Repo 🧑🏽‍💻
+https://github.com/salesforce/lavis
+
+
+## Additional Resources 📚
+- [BLIP-2 Demo Space](https://huggingface.co/spaces/hysts/BLIP2-with-transformers)
+- [BLIP-2 Transformers Example Notebooks](https://github.com/NielsRogge/Transformers-Tutorials/tree/master/BLIP-2) by Niels Rogge
+- [BLIP-2 Transformers Docs](https://huggingface.co/docs/transformers/model_doc/blip-2)

computer-vision-study-group/Sessions/Fiber.md

@@ -0,0 +1,24 @@
+# Fiber: Coarse-to-Fine Vision-Language Pre-Training with Fusion in the Backbone
+Session by [johko](https://github.com/johko)
+
+
+## Recording 📺
+[YouTube](https://www.youtube.com/watch?v=m9qhNGuWE2g&t=20s&pp=ygUdaHVnZ2luZyBmYWNlIHN0dWR5IGdyb3VwIHN3aW4%3D)
+
+
+## Session Slides 🖥️
+[Google Drive](https://docs.google.com/presentation/d/1vSu27tE87ZM103_CkgqsW7JeIp2mrmyl/edit?usp=sharing&ouid=107717747412022342990&rtpof=true&sd=true)
+
+
+## Original Paper 📄
+[Hugging Face](https://huggingface.co/papers/2206.07643) /
+[arxiv](https://arxiv.org/abs/2206.07643)
+
+
+## GitHub Repo 🧑🏽‍💻
+https://github.com/microsoft/fiber
+
+
+## Additional Resources 📚
+- [Text to Pokemon](https://huggingface.co/spaces/lambdalabs/text-to-pokemon) HF Space to create your own Pokemon
+- [Paper to Pokemon](https://huggingface.co/spaces/hugging-fellows/paper-to-pokemon) derived from the above space - create your own Pokemon from a paper

computer-vision-study-group/Sessions/FlexiViT.md

@@ -0,0 +1,23 @@
+# FlexiViT: One Model for All Patch Sizes
+Session by [johko](https://github.com/johko)
+
+
+## Recording 📺
+[YouTube](https://www.youtube.com/watch?v=TlRYBgsl7Q8&t=977s&pp=ygUdaHVnZ2luZyBmYWNlIHN0dWR5IGdyb3VwIHN3aW4%3D)
+
+
+## Session Slides 🖥️
+[Google Drive](https://docs.google.com/presentation/d/1rLAYr160COYQMUN0FDH7D9pP8qe1_QyXGvfbHkutOt8/edit?usp=sharing)
+
+
+## Original Paper 📄
+[Hugging Face](https://huggingface.co/papers/2212.08013) /
+[arxiv](https://arxiv.org/abs/2212.08013)
+
+
+## GitHub Repo 🧑🏽‍💻
+https://github.com/google-research/big_vision
+
+
+## Additional Resources 📚
+- [FlexiViT PR](https://github.com/google-research/big_vision/pull/24)

computer-vision-study-group/Sessions/HFVisionEcosystem.md

@@ -0,0 +1,10 @@
+# Hugging Face Vision Ecosystem Overview (June 2022)
+Session by [Niels Rogge](https://github.com/NielsRogge)
+
+
+## Recording 📺
+[YouTube](https://www.youtube.com/watch?v=oL-xmufhZM8&pp=ygUdaHVnZ2luZyBmYWNlIHN0dWR5IGdyb3VwIHN3aW4%3D)
+
+
+## Additional Resources 📚
+- [Accompanying Notebook](../Notebooks/HuggingFace_vision_ecosystem_overview_(June_2022).ipynb)

computer-vision-study-group/Sessions/HowDoVisionTransformersWork.md

@@ -0,0 +1,27 @@
+# How Do Vision Transformers Work
+Session by [johko](https://github.com/johko)
+
+
+## Session Slides 🖥️
+[Google Drive](https://docs.google.com/presentation/d/1PewOHVABkxx0jO9PoJSQi8to_WNlL4HdDp4M9e4L8hs/edit?usp=drivesdks)
+
+
+## Original Paper 📄
+[Hugging Face](https://huggingface.co/papers/2202.06709) /
+[arxiv](https://arxiv.org/pdf/2202.06709.pdf)
+
+
+## GitHub Repo 🧑🏽‍💻
+https://github.com/microsoft/Swin-Transformer
+
+
+## Additional Resources 📚
+Hessian Matrices:
+
+- https://stackoverflow.com/questions/23297090/how-calculating-hessian-works-for-neural-network-learning
+- https://machinelearningmastery.com/a-gentle-introduction-to-hessian-matrices/
+
+Loss Landscape Visualization:
+
+- https://mathformachines.com/posts/visualizing-the-loss-landscape/
+- https://github.com/tomgoldstein/loss-landscape

computer-vision-study-group/Sessions/MaskedAutoEncoders.md

@@ -0,0 +1,24 @@
+# Masked Autoencoders are Scalable Vision Learners
+Session by [johko](https://github.com/johko)
+
+
+## Recording 📺
+[YouTube](https://www.youtube.com/watch?v=AC6flxUFLrg&pp=ygUdaHVnZ2luZyBmYWNlIHN0dWR5IGdyb3VwIHN3aW4%3D)
+
+
+## Session Slides 🖥️
+[Google Drive](https://docs.google.com/presentation/d/10ZZ-Rl1D57VX005a58OmqNeOB6gPnE54/edit?usp=sharing&ouid=107717747412022342990&rtpof=true&sd=true)
+
+
+## Original Paper 📄
+[Hugging Face](https://huggingface.co/papers/2111.06377) /
+[arxiv](https://arxiv.org/abs/2111.06377)
+
+
+## GitHub Repo 🧑🏽‍💻
+https://github.com/facebookresearch/mae
+
+
+## Additional Resources 📚
+- [Transformers Docs ViTMAE](https://huggingface.co/docs/transformers/model_doc/vit_mae)
+- [Transformers ViTMAE Demo Notebook](https://github.com/NielsRogge/Transformers-Tutorials/tree/master/ViTMAE) by Niels Rogge

computer-vision-study-group/Sessions/NeuralRadianceFields.md

@@ -0,0 +1,19 @@
+# Introduction to Neural Radiance Fields
+Session by [Aritra](https://arig23498.github.io/) and [Ritwik](ritwikraha.github.io)
+
+
+## Recording 📺
+[YouTube](https://www.youtube.com/watch?v=U2XS7SxOy2s)
+
+
+## Session Slides 🖥️
+[Google Drive](https://docs.google.com/presentation/d/e/2PACX-1vTQVnoTJGhRxDscNV1Mg2aYhvXP8cKODpB5Ii72NWoetCGrTLBJWx_UD1oPXHrzPtj7xO8MS_3TQaSH/pub?start=false&loop=false&delayms=3000)
+
+
+## Original Paper 📄
+[Hugging Face](https://huggingface.co/papers/2003.08934) /
+[arxiv](https://arxiv.org/abs/2003.08934)
+
+
+## GitHub Repo 🧑🏽‍💻
+https://github.com/bmild/nerf

computer-vision-study-group/Sessions/PolarizedSelfAttention.md

@@ -0,0 +1,14 @@
+# Polarized Self-Attention
+Session by [Satpal](https://github.com/satpalsr)
+
+## Session Slides 🖥️
+[GitHub PDF](https://github.com/satpalsr/Talks/blob/main/PSA_discussion.pdf)
+
+
+## Original Paper 📄
+[Hugging Face](https://huggingface.co/papers/2107.00782) /
+[arxiv](https://arxiv.org/pdf/2107.00782.pdf)
+
+
+## GitHub Repo 🧑🏽‍💻
+https://github.com/DeLightCMU/PSA

computer-vision-study-group/Sessions/SwinTransformer.md

@@ -0,0 +1,25 @@
+# Swin Transformer
+Session by [johko](https://github.com/johko)
+
+
+## Recording 📺
+[YouTube](https://www.youtube.com/watch?v=Ngikt-K1Ecc&t=305s&pp=ygUdaHVnZ2luZyBmYWNlIHN0dWR5IGdyb3VwIHN3aW4%3D)
+
+
+## Session Slides 🖥️
+[Google Drive](https://docs.google.com/presentation/d/1RoFIC6vE55RS4WNqSlzNu3ljB6F-_8edtprAFXpGvKs/edit?usp=sharing)
+
+
+## Original Paper 📄
+[Hugging Face](https://huggingface.co/papers/2103.14030) /
+[arxiv](https://arxiv.org/pdf/2103.14030.pdf)
+
+
+## GitHub Repo 🧑🏽‍💻
+https://github.com/xxxnell/how-do-vits-work
+
+
+## Additional Resources 📚
+- [Transformers Docs Swin v1](https://huggingface.co/docs/transformers/model_doc/swin)
+- [Transformers Docs Swin v2](https://huggingface.co/docs/transformers/model_doc/swinv2)
+- [Transformers Docs Swin Super Resolution](https://huggingface.co/docs/transformers/model_doc/swin2sr)

config.json

@@ -0,0 +1,52 @@
+{
+  "_name_or_path": "openai/whisper-small",
+  "activation_dropout": 0.0,
+  "activation_function": "gelu",
+  "apply_spec_augment": false,
+  "architectures": [
+    "WhisperForConditionalGeneration"
+  ],
+  "attention_dropout": 0.0,
+  "begin_suppress_tokens": [
+    220,
+    50257
+  ],
+  "bos_token_id": 50257,
+  "classifier_proj_size": 256,
+  "d_model": 768,
+  "decoder_attention_heads": 12,
+  "decoder_ffn_dim": 3072,
+  "decoder_layerdrop": 0.0,
+  "decoder_layers": 12,
+  "decoder_start_token_id": 50258,
+  "dropout": 0.0,
+  "encoder_attention_heads": 12,
+  "encoder_ffn_dim": 3072,
+  "encoder_layerdrop": 0.0,
+  "encoder_layers": 12,
+  "eos_token_id": 50257,
+  "forced_decoder_ids": null,
+  "init_std": 0.02,
+  "is_encoder_decoder": true,
+  "mask_feature_length": 10,
+  "mask_feature_min_masks": 0,
+  "mask_feature_prob": 0.0,
+  "mask_time_length": 10,
+  "mask_time_min_masks": 2,
+  "mask_time_prob": 0.05,
+  "max_length": 448,
+  "max_source_positions": 1500,
+  "max_target_positions": 448,
+  "median_filter_width": 7,
+  "model_type": "whisper",
+  "num_hidden_layers": 12,
+  "num_mel_bins": 80,
+  "pad_token_id": 50257,
+  "scale_embedding": false,
+  "suppress_tokens": [],
+  "torch_dtype": "float32",
+  "transformers_version": "4.40.0.dev0",
+  "use_cache": false,
+  "use_weighted_layer_sum": false,
+  "vocab_size": 51865
+}

gradio-blocks/README.md

@@ -0,0 +1,123 @@
+# Welcome to the [Gradio](https://gradio.app/) Blocks Party 🥳
+
+![image (1)](https://user-images.githubusercontent.com/81195143/167954125-9854bf6b-4ae5-4735-8fdd-830fec41efa1.png)
+
+
+_**Timeline**: May 17th, 2022 - May 31st, 2022_
+
+---
+
+We are happy to invite you to the Gradio Blocks Party - a community event in which we will create **interactive demos** for state-of-the-art machine learning models. Demos are powerful because they allow anyone — not just ML engineers — to try out models in the browser, give feedback on predictions, identify trustworthy models. The event will take place from **May 17th to 31st**. We will be organizing this event on [Github](https://github.com/huggingface/community-events) and the [Hugging Face discord channel](https://discord.com/invite/feTf9x3ZSB). Prizes will be given at the end of the event, see: [Prizes](#prizes)
+
+<img src="https://user-images.githubusercontent.com/81195143/168656398-ace7acc9-ef7a-4e90-a9cd-c7d15dd800e1.gif" width="1160" height="600"/>
+
+## What is Gradio?
+
+Gradio is a Python library that allows you to quickly build web-based machine learning demos, data science dashboards, or other kinds of web apps, entirely in Python. These web apps can be launched from wherever you use Python (jupyter notebooks, colab notebooks, Python terminal, etc.) and shared with anyone instantly using Gradio's auto-generated share links. To learn more about Gradio see the Getting Started Guide: https://gradio.app/getting_started/ and the new Course on Huggingface about Gradio: [Gradio Course](https://huggingface.co/course/chapter9/1?fw=pt).
+
+Gradio can be installed via pip and comes preinstalled in Hugging Face Spaces, the latest version of Gradio can be set in the README in spaces by setting the sdk_version for example `sdk_version: 3.0b8`
+
+`pip install gradio` to install gradio locally
+
+
+## What is Blocks?
+
+`gradio.Blocks` is a low-level API that allows you to have full control over the data flows and layout of your application. You can build very complex, multi-step applications using Blocks. If you have already used `gradio.Interface`, you know that you can easily create fully-fledged machine learning demos with just a few lines of code. The Interface API is very convenient but in some cases may not be sufficiently flexible for your needs. For example, you might want to:
+
+* Group together related demos as multiple tabs in one web app.
+* Change the layout of your demo instead of just having all of the inputs on the left and outputs on the right.
+* Have multi-step interfaces, in which the output of one model becomes the input to the next model, or have more flexible data flows in general.
+* Change a component's properties (for example, the choices in a Dropdown) or its visibility based on user input.
+
+To learn more about Blocks, see the [official guide](https://www.gradio.app/introduction_to_blocks/) and the [docs](https://gradio.app/docs/).
+
+## What is Hugging Face Spaces?
+
+Spaces are a simple way to host ML demo apps directly on your profile or your organization’s profile on Hugging Face. This allows you to create your ML portfolio, showcase your projects at conferences or to stakeholders, and work collaboratively with other people in the ML ecosystem. Learn more about Spaces in the [docs](https://huggingface.co/docs/hub/spaces).
+
+## How Do Gradio and Hugging Face work together?
+
+Hugging Face Spaces is a free hosting option for Gradio demos. Spaces comes with 3 SDK options: Gradio, Streamlit and Static HTML demos. Spaces can be public or private and the workflow is similar to github repos. There are over 2000+ Gradio spaces currently on Hugging Face. Learn more about spaces and gradio: https://huggingface.co/docs/hub/spaces
+
+## Event Plan
+
+main components of the event consist of:
+
+1. Learning about Gradio and the new Blocks Feature
+2. Building your own Blocks demo using Gradio and Hugging Face Spaces
+3. Submitting your demo on Spaces to the Gradio Blocks Party Organization
+4. Share your blocks demo with a permanent shareable link 
+5. Win Prizes
+
+
+## Example spaces using Blocks
+
+<img width="1180" alt="mindseye-lite" src="https://user-images.githubusercontent.com/81195143/168619604-cf1ac733-c10e-487f-add4-8da48002dcff.png">
+
+- [dalle-mini](https://huggingface.co/spaces/dalle-mini/dalle-mini)([Code](https://huggingface.co/spaces/dalle-mini/dalle-mini/blob/main/app/gradio/app.py))
+- [mindseye-lite](https://huggingface.co/spaces/multimodalart/mindseye-lite)([Code](https://huggingface.co/spaces/multimodalart/mindseye-lite/blob/main/app.py))
+- [ArcaneGAN-blocks](https://huggingface.co/spaces/akhaliq/ArcaneGAN-blocks)([Code](https://huggingface.co/spaces/akhaliq/ArcaneGAN-blocks/blob/main/app.py))
+- [gr-blocks](https://huggingface.co/spaces/merve/gr-blocks)([Code](https://huggingface.co/spaces/merve/gr-blocks/blob/main/app.py))
+- [tortoisse-tts](https://huggingface.co/spaces/osanseviero/tortoisse-tts)([Code](https://huggingface.co/spaces/osanseviero/tortoisse-tts/blob/main/app.py))
+- [CaptchaCracker](https://huggingface.co/spaces/osanseviero/tortoisse-tts)([Code](https://huggingface.co/spaces/akhaliq/CaptchaCracker/blob/main/app.py))
+
+
+## To participate in the event
+
+- Join the organization for Blocks event
+    - [https://huggingface.co/Gradio-Blocks](https://huggingface.co/Gradio-Blocks)
+- Join the discord
+    - [discord](https://discord.com/invite/feTf9x3ZSB)
+
+
+Participants will be building and sharing Gradio demos using the Blocks feature. We will share a list of ideas of spaces that can be created using blocks or participants are free to try out their own ideas. At the end of the event, spaces will be evaluated and prizes will be given.
+
+
+## Potential ideas for creating spaces:
+
+
+- Trending papers from https://paperswithcode.com/
+- Models from huggingface model hub: https://huggingface.co/models
+- Models from other model hubs
+    - Tensorflow Hub: see example Gradio demos at https://huggingface.co/tensorflow
+    - Pytorch Hub: see example Gradio demos at https://huggingface.co/pytorch
+    - ONNX model Hub: see example Gradio demos at https://huggingface.co/onnx
+    - PaddlePaddle Model Hub: see example Gradio demos at https://huggingface.co/PaddlePaddle
+- participant ideas, try out your own ideas
+
+
+## Prizes
+- 1st place winner based on likes
+    - [Hugging Face PRO subscription](https://huggingface.co/pricing) for 1 year
+    - Embedding your Gradio Blocks demo in the Gradio Blog
+- top 10 winners based on likes
+    - Swag from [Hugging Face merch shop](https://huggingface.myshopify.com/): t-shirts, hoodies, mugs of your choice
+- top 25 winners based on likes
+    - [Hugging Face PRO subscription](https://huggingface.co/pricing) for 1 month
+- Blocks event badge on HF for all participants!
+
+## Prizes Criteria
+
+- Staff Picks
+- Most liked Spaces
+- Community Pick (voting)
+- Most Creative Space (voting)
+- Most Educational Space (voting)
+- CEO's pick (one prize for a particularly impactful demo), picked by @clem
+- CTO's pick (one prize for a particularly technically impressive demo), picked by @julien
+
+
+## Creating a Gradio demo on Hugging Face Spaces
+
+Once a model has been picked from the choices above or feel free to try your own idea, you can share a model in a Space using Gradio
+
+Read more about how to add [Gradio spaces](https://huggingface.co/blog/gradio-spaces).
+ 
+Steps to add Gradio Spaces to the Gradio Blocks Party org
+1. Create an account on Hugging Face
+2. Join the Gradio Blocks Party Organization by clicking "Join Organization" button in the organization page or using the shared link above
+3. Once your request is approved, add your space using the Gradio SDK and share the link with the community!
+
+## LeaderBoard for Most Popular Blocks Event Spaces based on Likes
+
+- See Leaderboard: https://huggingface.co/spaces/Gradio-Blocks/Leaderboard

huggan/README.md

@@ -0,0 +1,487 @@
+# HugGAN Sprint
+
+![Banner](assets/huggan_banner.png?raw=true "Banner")
+
+_**Timeline**: April 4th, 2022 - April 17th, 2022_
+
+---
+
+Welcome to HugGAN Sprint! The goal of this sprint is to add more GANs and GAN-based demos to the Hugging Face Hub 🤗.
+
+During the sprint, we’ll be bringing in some awesome speakers to talk about GANs and the future of generative models. Oh, and if you need access to compute for your project, we’ll help you there too! As an added bonus, if you choose to participate, we’ll send you a gift (specific details TBD). We encourage you to form teams of ~2-3 people! Make friends in the Discord :)
+
+To join:
+
+1. Fill out [this form](https://forms.gle/goq41UgzsvuKKTFFA), so we can keep track of who’s joining.
+2. Send a reaction in the [#join-sprint channel](https://discord.com/channels/879548962464493619/954070850645135462) under the HugGAN category in Discord. This will add you to the rest of the related channels. If you haven't joined our discord yet, [click here](https://discord.gg/H3bUrDPTfS).
+3. Once you’ve decided what you want to work on, add your project’s information to [this sheet](https://docs.google.com/spreadsheets/d/1aAHqOOk2SOw4j6mrJLkLT6ZyKyLDOvGF5D9tuUqnoG8/edit#gid=0), where you can describe your project and let us know if you need additional compute. Still brainstorming? Feel free to propose ideas in #sprint-discussions.
+
+## Table of Contents
+
+- [Important dates](#important-dates)
+- [How to install relevant libraries](#how-to-install-relevant-libraries)
+- [General workflow](#general-workflow)
+- [Datasets to add](#datasets-to-add)
+- [Links to check out](#links-to-check-out)
+- [GAN metrics](#gan-metrics)
+- [Evaluation](#evaluation)
+- [Prizes](#prizes)
+- [Communication and Problems](#communication-and-problems)
+- [Talks](#talks)
+- [General Tips & Tricks](#general-tips-and-tricks)
+
+## Important dates
+
+| Date      | Description |
+| ----------- | ----------- |
+| April 4th      | Sprint Kickoff 🚀      |
+| April 15th   | Submission Deadline 🛑  |
+| April 22nd | Prizes Announced for Participants 🎁 |
+
+## How to install relevant libraries
+
+You'll need the following dependencies installed to use this repo:
+
+- [PyTorch](https://pytorch.org/) or [Keras](https://keras.io/) - depending on which framework you prefer ;)
+- [🤗 Datasets](https://huggingface.co/docs/datasets/index)
+- [🤗 Accelerate](https://huggingface.co/docs/accelerate/index) - in case you're planning to train a PyTorch model and you want it to be run effortlessly 
+
+We recommend installing the above libraries in a [virtual environment](https://docs.python.org/3/library/venv.html). 
+If you're unfamiliar with Python virtual environments, check out the [user guide](https://packaging.python.org/guides/installing-using-pip-and-virtual-environments/). Create a virtual environment with the version of Python you're going to use and activate it.
+
+You should be able to run the command:
+
+```bash
+python3 -m venv <your-venv-name>
+```
+
+You can activate your venv by running
+
+```bash
+source ~/<your-venv-name>/bin/activate
+```
+
+### Install Dependencies
+
+We've packaged up the example scripts here into a simple Python package. To install it, just pip install it
+
+```
+git clone https://github.com/huggingface/community-events.git
+cd community-events
+pip install .
+```
+
+If you use `pip install -e .` instead of `pip install`, it will install the package in development mode, which can be useful if you are planning on contributing any changes here 🤗. 
+
+## General workflow
+
+The process to follow is outlined below. It consists of 3 steps:
+
+1. Get a dataset and push to the Hub
+2. Train a model and push to the Hub
+3. Create a demo (🤗 Space)
+
+These steps are explained in more detail below.
+
+### 1. Get a dataset and push to Hub
+
+The first step is the most obvious one: to train a GAN (or any neural network), we need a dataset. This could be either a dataset that is already available on the [Hub](https://huggingface.co/datasets), or one that isn't already. Below we'll explain how to load the data in both cases.
+
+Note that we maintain a list of interesting datasets to add to the Hub [here](#datasets-to-add).
+
+#### 1.1 Use a dataset already available on the Hub
+
+Most famous computer vision datasets are already available on the [Hub](https://huggingface.co/datasets?task_categories=task_categories:image-classification) (such as [MNIST](https://huggingface.co/datasets/mnist), [Fashion MNIST](https://huggingface.co/datasets/fashion_mnist), [CIFAR-10](https://huggingface.co/datasets/cifar10), [CIFAR-100](https://huggingface.co/datasets/cifar100), etc.).
+
+Loading a dataset can be done as follows:
+
+```python
+from datasets import load_dataset
+
+# a general one ...
+dataset = load_dataset("mnist")
+
+# ... or one that's part of the huggan organization
+dataset = load_dataset("huggan/edges2shoes")
+```
+
+In a notebook, you can **directly see** the images by selecting a split and then the appropriate column:
+
+```python
+example = dataset['train'][0]
+print(example['image'])
+```
+
+#### 1.2 Upload a new dataset to the Hub
+
+In case your dataset is not already on the Hub, you can upload it to the `huggan` [organization](https://huggingface.co/huggan). If you've signed up for the event by filling in the [spreadsheet]((https://docs.google.com/spreadsheets/d/1aAHqOOk2SOw4j6mrJLkLT6ZyKyLDOvGF5D9tuUqnoG8/edit#gid=0)), your Hugging Face account should be part of it. 
+
+Let's illustrate with an example how this was done for NVIDIA's [MetFaces dataset](https://github.com/NVlabs/metfaces-dataset):
+
+<p align="center">
+    <img src="https://github.com/NVlabs/metfaces-dataset/blob/master/img/metfaces-teaser.png" alt="drawing" width="700"/>
+</p>
+
+Previously, this dataset was only hosted on [Google Drive](https://github.com/NVlabs/metfaces-dataset#overview), and not really easily accessible.
+
+To begin with, one should check that one is correctly logged in and that `git-lfs` is installed so that the dataset can be uploaded.
+
+Run:
+
+```bash
+huggingface-cli login
+```
+
+in a terminal, or case you're working in a notebook
+
+```python
+from huggingface_hub import notebook_login
+
+notebook_login()
+```
+
+It is recommended to login with your access token that can be found under your HuggingFace profile (icon in the top right corner on [hf.co](http://hf.co/), then Settings -> Access Tokens -> User Access Tokens -> New Token (if you haven't generated one already). Alternatively, you can go to [your token settings](https://huggingface.co/settings/tokens) directly.
+
+You can then copy-paste this token to log in locally.
+
+Next, let's make sure that `git-lfs` is correctly installed. To so, simply run:
+
+```bash
+git-lfs -v
+```
+
+The output should show something like `git-lfs/2.13.2 (GitHub; linux amd64; go 1.15.4)`. If your console states that the `git-lfs` command was not found, please make sure to install it [here](https://git-lfs.github.com/) or simply via: 
+
+```bash
+sudo apt-get install git-lfs
+git config --global user.email "you@example.com"
+git config --global user.name "Your Name"
+```
+
+Next, one can leverage the [`ImageFolder`](https://huggingface.co/docs/datasets/v2.0.0/en/image_process#imagefolder) builder to very easily upload an image dataset to the hub. In case the dataset you're uploading has a direct download URL, you can simply provide it to the `data_files` argument as shown below. Otherwise, you'll need to go to the link of the dataset and manually download it first as a zip/tar (which was the case for MetFaces), and provide the file through the `data_files` argument. Alternatively, it may be that you have a folder with images, in which case you can provide it using the `data_dir` argument. Note that the latter assumes a [particular structure](https://huggingface.co/docs/datasets/v2.0.0/en/image_process#imagefolder).
+
+```python
+from datasets import load_dataset
+
+# option 1: local folder
+dataset = load_dataset("imagefolder", data_dir="path_to_folder")
+# option 2: local or remote file(s), supporting the following extensions: tar, gzip, zip, xz, rar, zstd
+dataset = load_dataset("imagefolder", data_files="path_to_file_or_direct_download_link")
+
+# note that you can also provide them as separate splits, like so:
+dataset = load_dataset("imagefolder", data_files={"train": ["path/to/file1", "path/to/file2"], "test": ["path/to/file3", "path/to/file4"]})
+```
+
+Once you've loaded your dataset, you can push it to the Hub with a single line of code:
+
+```python
+dataset.push_to_hub("huggan/name-of-your-dataset")
+```
+
+Et voila! Your dataset is now available on the Hub :) If you wait a bit, the Dataset viewer should be able to preview images in the browser. The MetFaces dataset can be seen here: https://huggingface.co/datasets/huggan/metfaces. 
+
+<p align="center">
+    <img src="https://github.com/huggingface/community-events/blob/main/huggan/assets/metfaces.png" alt="drawing" width="700"/>
+</p>
+
+The cool thing is that anyone can now access this dataset from anywhere, using `load_dataset` 🎉🥳 this means that you can easily load the dataset on another computer for instance, or in a different environment. Amazing, isn't it?
+
+❗ Note: When uploading a dataset, make sure that it has appropriate column names. The `ImageFolder` utility automatically creates `image` and `label` columns, however if there's only one image class, it makes sense to remove the `label` column before pushing to the hub. This can be done as follows:
+
+```python
+dataset = dataset.remove_columns("label")
+```
+
+Note that you can always update a dataset by simply calling `push_to_hub` again (providing the same name).
+
+#### 1.3 Processing the data
+
+Once you've uploaded your dataset, you can load it and create a dataloader for it. The code example below shows how to apply some data augmentation and creating a PyTorch Dataloader (the [PyTorch example scripts](pytorch) all leverage this). More info can also be found in the [docs](https://huggingface.co/docs/datasets/v2.0.0/en/image_process#process-image-data).
+
+```python
+from datasets import load_dataset
+from torchvision.transforms import CenterCrop, Compose, Normalize, Resize, ToTensor
+from torch.utils.data import DataLoader
+
+# load your data
+dataset = load_dataset("dataset_name")
+
+image_size = 256
+
+# define image transformations (e.g. using torchvision)
+transform = Compose(
+    [
+        Resize(image_size),
+        CenterCrop(image_size),
+        ToTensor(),
+        Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
+    ]
+)
+
+# define function
+def transforms(examples):
+   examples["image"] = [transform(image.convert("RGB")) for image in examples["image"]]
+
+   return examples
+
+transformed_dataset = dataset.with_transform(transforms)
+
+# create dataloader
+dataloader = DataLoader(
+     transformed_dataset["train"], batch_size="your batch size", shuffle=True, num_workers="your number of CPU cores"
+)
+``` 
+
+As can be seen, we leverage the [`with_transform`](https://huggingface.co/docs/datasets/v2.0.0/en/package_reference/main_classes#datasets.Dataset.with_transform) method here, which will make sure the image transformations will only be performed when iterating over the data (i.e. data augmentation is performed on-the-fly, making it very RAM-friendly) rather than performing it on the entire dataset in one go (which would be the case if you use [`map`](https://huggingface.co/docs/datasets/v2.0.0/en/package_reference/main_classes#datasets.Dataset.map)). The `with_transform` method does the same thing as [`set_transform`](https://huggingface.co/docs/datasets/v2.0.0/en/package_reference/main_classes#datasets.Dataset.set_transform), except that it does return a new `Dataset` rather than performing the operation in-place.
+
+### 2. Train a model and push to Hub
+
+Next, one can start training a model. This could be any model you'd like. However, we provide some example scripts to help you get started, in both [PyTorch](pytorch) and [Tensorflow](tensorflow). An example is the [DCGAN](pytorch/dcgan) model for unconditional image generation. Simply follow the README that explains all the details of the relevant implementation, and run it in your environment. 
+
+The PyTorch example scripts all leverage 🤗 [Accelerate](https://huggingface.co/docs/accelerate/index), which provides an easy API to make your scripts run on any kind of distributed setting (multi-GPUs, TPUs etc.) and with mixed precision, while still letting you write your own training loop. 
+
+Alternatively, we also provide a [Links to Check Out](#links-to-check-out) section to give you some inspiration.
+
+Below, we explain in more detail how to upload your model to the Hub, depending on the framework you're using (sections [2.1](#21-pytorch) and [2.2](#22-keras)). In section [2.3](#33-alternative-ways-to-upload-a-model-to-the-hub), we'll explain how to write a nice model card. In section [2.4](24-model-cards), we'll illustrate alternative ways to upload (and re-use) a model to (and from) the hub. Finally, in section [2.5](25-accelerate), we explain 🤗 [Accelerate](https://huggingface.co/docs/accelerate/index), the awesome library that makes training PyTorch models on any kind of environment a breeze. Be sure to check it out!
+
+#### 2.1 PyTorch
+
+If you're planning to train a custom PyTorch model, it's recommended to make it inherit from `PyTorchModelHubMixin`. This makes sure you can push it to the Hub at the end of training, and reload it afterwards using `from_pretrained`, as shown in the code example below:
+
+```python
+from huggingface_hub import PyTorchModelHubMixin
+
+class MyGenerator(nn.Module, PyTorchModelHubMixin):
+   def __init__(self, **kwargs):
+      super().__init__()
+      self.config = kwargs.pop("config", None)
+      self.layer = ...
+   def forward(self, ...):
+      return ...
+
+# Create model
+model = MyGenerator()
+
+# Push to HuggingFace Hub
+model.push_to_hub("huggan/name-of-your-model").
+
+# Reload from HuggingFace Hub
+reloaded = MyGenerator.from_pretrained("huggan/name-of-your-model").
+```
+
+This `PyTorchModelHubMixin` class is available in the [`huggingface_hub` library](https://github.com/huggingface/huggingface_hub), which comes pre-installed if you install `datasets` (or `transformers`) in your environment.
+
+#### 2.2 Keras
+
+In Keras, one can leverage the `push_to_hub_keras` and `from_pretrained_keras` methods:
+
+```python
+import tensorflow as tf
+from huggingface_hub import push_to_hub_keras, from_pretrained_keras
+
+# Build a Keras model
+inputs = tf.keras.layers.Input(shape=(2,))
+x = tf.keras.layers.Dense(2, activation="relu")(inputs)
+model = tf.keras.models.Model(inputs=inputs, outputs=x)
+model.compile(optimizer="adam", loss="mse")
+
+# Push to HuggingFace Hub
+push_to_hub_keras(model, "huggan/my-cool-model")
+
+# Reload from HuggingFace Hub
+reloaded = from_pretrained_keras("huggan/my-cool-model")
+```
+
+These methods are available in the [`huggingface_hub` library](https://github.com/huggingface/huggingface_hub), which comes pre-installed if you install `datasets` (or `transformers`) in your environment. Note that the `push_to_hub_keras` method supports pushing several models (such as a generator and discriminator) to the same repo, as illustrated [here](https://github.com/huggingface/huggingface_hub/issues/533#issuecomment-1058093158).
+
+#### 2.3 Alternative ways to upload a model to the Hub
+
+Besides the methods explained in sections 2.1 and 2.2 above, you can also share model assets directly from git, which is explained in depth in [this guide](https://huggingface.co/docs/hub/adding-a-model#uploading-your-files).
+
+#### 2.4 Model cards
+
+When uploading a model to the Hub, it's important to include a so-called [model card](https://huggingface.co/course/chapter4/4?fw=pt) with it. This is just a README (in Markdown) 🃏 that includes:
+- license,
+- task,
+- `huggan` and `gan` tags,
+- dataset metadata,
+- information related to the model,
+- information on dataset, intended uses,
+- a model output.
+
+If you trained one of the example models, this model card will be automatically generated for you. If you didn’t train the model yourself, be sure to both credit the original authors and include the associated license in your model card! Here is an [example model repo](https://huggingface.co/merve/anime-faces-generator).
+
+You can also use this [template model card](model_card_template.md)
+ as a guide to build your own.
+
+![Alt text](assets/example_model.png?raw=true "Title")
+
+#### 2.5 Accelerate
+
+HuggingFace `accelerate` is an awesome library for training PyTorch models. Here we show why. 
+
+Basically, the library requires to replace this:
+
+```
+my_model.to(device)
+
+for batch in my_training_dataloader:
+    my_optimizer.zero_grad()
+    inputs, targets = batch
+    inputs = inputs.to(device)
+    targets = targets.to(device)
+    outputs = my_model(inputs)
+    loss = my_loss_function(outputs, targets)
+    loss.backward()
+    my_optimizer.step()
+```    
+    
+by this:
+
+```diff
++ from accelerate import Accelerator
+
++ accelerator = Accelerator()
+- my_model.to(device)
+  # Pass every important object (model, optimizer, dataloader) to *accelerator.prepare*
++ my_model, my_optimizer, my_training_dataloader = accelerate.prepare(
++     my_model, my_optimizer, my_training_dataloader
++ )
+
+  for batch in my_training_dataloader:
+      my_optimizer.zero_grad()
+      inputs, targets = batch
+-     inputs = inputs.to(device)
+-     targets = targets.to(device)
+      outputs = my_model(inputs)
+      loss = my_loss_function(outputs, targets)
+      # Just a small change for the backward instruction
+-     loss.backward()
++     accelerator.backward(loss)
+      my_optimizer.step()
+```
+
+and BOOM, your script runs on **any kind of hardware**, including CPU, multi-CPU, GPU, multi-GPU and TPU. It also supports things like [DeepSpeed](https://github.com/microsoft/DeepSpeed) and [mixed precision](https://arxiv.org/abs/1710.03740) for training efficiently.
+
+You can now run your script as follows:
+
+```bash
+accelerate config
+```
+
+=> Accelerate will ask what kind of environment you'd like to run your script on, simply answer the questions being asked. Next:
+
+```bash
+accelerate launch <your script.py>
+```
+
+This will run your script on the environment you asked for. You can always check the environment settings by typing:
+
+```bash
+accelerate env
+```
+
+You can of course change the environment by running `accelerate config` again.
+
+### 3. Create a demo
+
+Once you share a model, you then should share a [Space](https://huggingface.co/spaces) based on your SDK of choice (Gradio or Streamlit) or as a static page. 🌌
+
+![Alt text](assets/example_space.png?raw=true "Title")
+
+Here is an [example Space](https://huggingface.co/spaces/merve/anime-face-generator) corresponding to the model example shared above. Don’t know how to create a space? Read more about how to add spaces [here](https://huggingface.co/docs/hub/spaces).
+
+Below, we list some other great example GAN Spaces:
+- AnimeGANv2: https://huggingface.co/spaces/akhaliq/AnimeGANv2
+- ArcaneGAN: https://huggingface.co/spaces/akhaliq/ArcaneGAN
+- This Pokemon does not exist: https://huggingface.co/spaces/ronvolutional/ai-pokemon-card
+- GFP-GAN: https://huggingface.co/spaces/akhaliq/GFPGAN
+- DualStyleGAN: https://huggingface.co/spaces/hysts/DualStyleGAN
+
+## Example Scripts
+
+In this repo, we have provided some example scripts you can use to train your own GANs. Below is a table of the available scripts:
+
+| Name      | Paper |
+| ----------- | ----------- |
+| [DCGAN](pytorch/dcgan)  | [Unsupervised Representation Learning with Deep Convolutional Generative Adversarial Networks](https://arxiv.org/abs/1511.06434)  |
+| [pix2pix](pytorch/pix2pix) | [Image-to-Image Translation with Conditional Adversarial Networks](https://arxiv.org/abs/1611.07004) |
+| [CycleGAN](pytorch/cyclegan)  | [Unpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networks](https://arxiv.org/abs/1703.10593)
+
+## Datasets to add
+
+Below, we list some datasets which could be added to the Hub (feel free to add on one of these, or open a PR to add more datasets!):
+
+- DeepFashion: https://mmlab.ie.cuhk.edu.hk/projects/DeepFashion.html
+- Flowers: https://www.robots.ox.ac.uk/~vgg/data/flowers/
+- LSUN: https://www.yf.io/p/lsun
+
+## Links to Check Out
+
+Below, we list some possible awesome project ideas (feel free to work on one of these, or open a PR to add more project ideas!):
+
+PyTorch:
+- Lightweight-GAN: https://github.com/lucidrains/lightweight-gan
+- StyleGAN2: https://github.com/lucidrains/stylegan2-pytorch
+- StyleGAN2-ada: https://github.com/NVlabs/stylegan2-ada
+- StyleGAN3 (alias-free GAN): https://github.com/NVlabs/stylegan3
+- BigGAN: https://github.com/ajbrock/BigGAN-PyTorch, https://github.com/huggingface/pytorch-pretrained-BigGAN
+- ADGAN: https://github.com/menyifang/ADGAN
+- ICGAN: https://github.com/facebookresearch/ic_gan
+- StarGANv2: https://github.com/clovaai/stargan-v2
+- Progressive Growing GAN: https://github.com/Maggiking/PGGAN-PyTorch
+- Vision Aided GAN: https://github.com/nupurkmr9/vision-aided-gan
+- DiffAugment (for training data-efficient GANs): https://github.com/mit-han-lab/data-efficient-gans
+- StyleGAN-XL: https://github.com/autonomousvision/stylegan_xl
+- CUT: https://github.com/taesungp/contrastive-unpaired-translation
+- studioGAN (library with many GAN implementations): https://github.com/POSTECH-CVLab/PyTorch-StudioGAN
+- MMGeneration (library with many GAN implementations): https://github.com/open-mmlab/mmgeneration
+- Deformable GAN: https://github.com/ssfootball04/pose-transfer
+- Denoising Diffusion GAN: https://github.com/NVlabs/denoising-diffusion-gan
+
+Keras:
+- WGAN-GP: https://keras.io/examples/generative/wgan_gp/
+- Conditional GAN: https://keras.io/examples/generative/conditional_gan/
+- CycleGAN, DiscoGAN etc.: https://github.com/eriklindernoren/Keras-GAN
+- Neural Style Transfer: https://www.tensorflow.org/tutorials/generative/style_transfer
+- Image Super Resolution: https://github.com/idealo/image-super-resolution
+- Deformable GAN: https://github.com/AliaksandrSiarohin/pose-gan
+
+General links & tutorials:
+- https://github.com/yhlleo/GAN-Metrics
+- https://paperswithcode.com/task/image-generation
+
+## GAN metrics
+
+There have been several quantitative measures defined for assessing the quality of GANs (and other generative models). Refer to [this page](pytorch/metrics) for more info.
+
+## Evaluation
+
+For each submission, you are expected to submit:
+
+1. A model repository
+2. A space made with the model repository you created
+
+## Prizes
+
+TODO
+
+## Communication and Problems
+
+If you encounter any problems or have any questions, you should use one of the following platforms depending on your type of problem. Hugging Face is an "open-source-first" organization meaning  that we'll try to solve all problems in the most public and most transparent way possible so that everybody in the community profits.
+
+The following table summarizes what platform to use for which problem.
+
+- Problem/question/bug with the 🤗 Datasets library that you think is a general problem that also impacts other people, please open an [Issues on Datasets](https://github.com/huggingface/datasets/issues/new?assignees=&labels=bug&template=bug-report.md&title=) and ping @nielsrogge.
+- Problem/question with a modified, customized training script that is less likely to impact other people, please post your problem/question [on the forum](https://discuss.huggingface.co/) and ping @nielsrogge.
+- Other questions regarding the event, rules of the event, or if you are not sure where to post your question, please ask in the Discord channel [**#sprint-discussions**](https://discord.com/channels/879548962464493619/954111918895943720).
+
+## Talks
+
+TODO
+
+## General Tips and Tricks
+
+- Memory efficient training:
+
+In case, you are getting out-of-memory errors on your GPU, we recommend to use  [bitsandbytes](https://github.com/facebookresearch/bitsandbytes) to replace the native memory-intensive Adam optimizer with the one of `bitsandbytes`. It can be used to both train the generator and the discriminator in case you're training a GAN.

huggan/__init__.py

@@ -0,0 +1,3 @@
+from pathlib import Path
+
+TEMPLATE_MODEL_CARD_PATH = Path(__file__).parent.absolute() / 'model_card_template.md'

huggan/model_card_template.md

@@ -0,0 +1,50 @@
+---
+tags:
+- huggan
+- gan
+# See a list of available tags here:
+# https://github.com/huggingface/hub-docs/blob/main/js/src/lib/interfaces/Types.ts#L12
+# task: unconditional-image-generation or conditional-image-generation or image-to-image
+license: mit
+---
+
+# MyModelName
+
+## Model description
+
+Describe the model here (what it does, what it's used for, etc.)
+
+## Intended uses & limitations
+
+#### How to use
+
+```python
+# You can include sample code which will be formatted
+```
+
+#### Limitations and bias
+
+Provide examples of latent issues and potential remediations.
+
+## Training data
+
+Describe the data you used to train the model.
+If you initialized it with pre-trained weights, add a link to the pre-trained model card or repository with description of the pre-training data.
+
+## Training procedure
+
+Preprocessing, hardware used, hyperparameters...
+
+## Eval results
+
+## Generated Images
+
+You can embed local or remote images using `![](...)`
+
+### BibTeX entry and citation info
+
+```bibtex
+@inproceedings{...,
+  year={2020}
+}
+```

huggan/pytorch/README.md

@@ -0,0 +1,19 @@
+# Example scripts (PyTorch)
+
+This directory contains a few example scripts that allow you to train famous GANs on your own data using a bit of 🤗 magic.
+
+More concretely, these scripts:
+- leverage 🤗 [Datasets](https://huggingface.co/docs/datasets/index) to load any image dataset from the hub (including your own, possibly private, dataset)
+- leverage 🤗 [Accelerate](https://huggingface.co/docs/accelerate/index) to instantly run the script on (multi-) CPU, (multi-) GPU, TPU environments, supporting fp16 and mixed precision as well as DeepSpeed
+- leverage 🤗 [Hub](https://huggingface.co/) to push the model to the hub at the end of training, allowing to easily create a demo for it afterwards
+
+Currently, it contains the following examples:
+
+| Name      | Paper |
+| ----------- | ----------- |
+| [DCGAN](dcgan)  | [Unsupervised Representation Learning with Deep Convolutional Generative Adversarial Networks](https://arxiv.org/abs/1511.06434)  |
+| [pix2pix](pix2pix) | [Image-to-Image Translation with Conditional Adversarial Networks](https://arxiv.org/abs/1611.07004) |
+| [CycleGAN](cyclegan)  | [Unpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networks](https://arxiv.org/abs/1703.10593)
+| [Lightweight GAN](lightweight_gan) | [Towards Faster and Stabilized GAN Training for High-fidelity Few-shot Image Synthesis](https://openreview.net/forum?id=1Fqg133qRaI)
+
+

huggan/pytorch/cyclegan/README.md

@@ -0,0 +1,81 @@
+# Training CycleGAN on your own data
+
+This folder contains a script to train [CycleGAN](https://arxiv.org/abs/1703.10593), leveraging the [Hugging Face](https://huggingface.co/) ecosystem for processing data and pushing the model to the Hub. 
+
+<p align="center">
+    <img src="https://camo.githubusercontent.com/16fa02525bf502bec1aac77a3eb5b96928b0f25d73f7d9dedcc041ba28c38751/68747470733a2f2f6a756e79616e7a2e6769746875622e696f2f4379636c6547414e2f696d616765732f7465617365725f686967685f7265732e6a7067" alt="drawing" width="700"/>
+</p>
+
+Example applications of CycleGAN. Taken from [this repo](https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix).
+
+The script leverages 🤗 Datasets for loading and processing data, and 🤗 Accelerate for instantly running on CPU, single, multi-GPUs or TPU, also supporting mixed precision.
+
+## Launching the script
+
+To train the model with the default parameters (200 epochs, 256x256 images, etc.) on [huggan/facades](https://huggingface.co/datasets/huggan/facades) on your environment, first run:
+
+```bash
+accelerate config
+```
+
+and answer the questions asked. Next, launch the script as follows: 
+
+```
+accelerate launch train.py
+```
+
+This will create local "images" and "saved_models" directories, containing generated images and saved checkpoints over the course of the training.
+
+To train on another dataset available on the hub, simply do:
+
+```
+accelerate launch train.py --dataset huggan/edges2shoes
+```
+
+Make sure to pick a dataset which has "imageA" and "imageB" columns defined. One can always tweak the script in case the column names are different.
+
+## Training on your own data
+
+You can of course also train on your own images. For this, one can leverage Datasets' [ImageFolder](https://huggingface.co/docs/datasets/v2.0.0/en/image_process#imagefolder). Make sure to authenticate with the hub first, by running the `huggingface-cli login` command in a terminal, or the following in case you're working in a notebook:
+
+```python
+from huggingface_hub import notebook_login
+
+notebook_login()
+```
+
+Next, run the following in a notebook/script:
+
+```python
+from datasets import load_dataset
+
+# first: load dataset
+# option 1: from local folder
+dataset = load_dataset("imagefolder", data_dir="path_to_folder")
+# option 2: from remote URL (e.g. a zip file)
+dataset = load_dataset("imagefolder", data_files="URL to .zip file")
+
+# next: push to the hub (assuming git-LFS is installed)
+dataset.push_to_hub("huggan/my-awesome-dataset")
+```
+
+You can then simply pass the name of the dataset to the script:
+
+```
+accelerate launch train.py --dataset huggan/my-awesome-dataset
+```
+
+## Pushing model to the Hub
+
+You can push your trained generator to the hub after training by specifying the `push_to_hub` flag. 
+Then, you can run the script as follows:
+
+```
+accelerate launch train.py --push_to_hub --model_name cyclegan-horse2zebra
+```
+
+This is made possible by making the generator inherit from `PyTorchModelHubMixin`available in the `huggingface_hub` library. 
+
+# Citation
+
+This repo is entirely based on Erik Linder-Norén's [PyTorch-GAN repo](https://github.com/eriklindernoren/PyTorch-GAN), but with added HuggingFace goodies.

huggan/pytorch/cyclegan/modeling_cyclegan.py

@@ -0,0 +1,108 @@
+import torch.nn as nn
+import torch.nn.functional as F
+import torch
+
+from huggan.pytorch.huggan_mixin import HugGANModelHubMixin
+
+
+##############################
+#           RESNET
+##############################
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, in_features):
+        super(ResidualBlock, self).__init__()
+
+        self.block = nn.Sequential(
+            nn.ReflectionPad2d(1),
+            nn.Conv2d(in_features, in_features, 3),
+            nn.InstanceNorm2d(in_features),
+            nn.ReLU(inplace=True),
+            nn.ReflectionPad2d(1),
+            nn.Conv2d(in_features, in_features, 3),
+            nn.InstanceNorm2d(in_features),
+        )
+
+    def forward(self, x):
+        return x + self.block(x)
+
+
+class GeneratorResNet(nn.Module, HugGANModelHubMixin):
+    def __init__(self, input_shape, num_residual_blocks):
+        super(GeneratorResNet, self).__init__()
+
+        channels = input_shape[0]
+
+        # Initial convolution block
+        out_features = 64
+        model = [
+            nn.ReflectionPad2d(channels),
+            nn.Conv2d(channels, out_features, 7),
+            nn.InstanceNorm2d(out_features),
+            nn.ReLU(inplace=True),
+        ]
+        in_features = out_features
+
+        # Downsampling
+        for _ in range(2):
+            out_features *= 2
+            model += [
+                nn.Conv2d(in_features, out_features, 3, stride=2, padding=1),
+                nn.InstanceNorm2d(out_features),
+                nn.ReLU(inplace=True),
+            ]
+            in_features = out_features
+
+        # Residual blocks
+        for _ in range(num_residual_blocks):
+            model += [ResidualBlock(out_features)]
+
+        # Upsampling
+        for _ in range(2):
+            out_features //= 2
+            model += [
+                nn.Upsample(scale_factor=2),
+                nn.Conv2d(in_features, out_features, 3, stride=1, padding=1),
+                nn.InstanceNorm2d(out_features),
+                nn.ReLU(inplace=True),
+            ]
+            in_features = out_features
+
+        # Output layer
+        model += [nn.ReflectionPad2d(channels), nn.Conv2d(out_features, channels, 7), nn.Tanh()]
+
+        self.model = nn.Sequential(*model)
+
+    def forward(self, x):
+        return self.model(x)
+
+
+##############################
+#        Discriminator
+##############################
+
+
+class Discriminator(nn.Module):
+    def __init__(self, channels):
+        super(Discriminator, self).__init__()
+
+        def discriminator_block(in_filters, out_filters, normalize=True):
+            """Returns downsampling layers of each discriminator block"""
+            layers = [nn.Conv2d(in_filters, out_filters, 4, stride=2, padding=1)]
+            if normalize:
+                layers.append(nn.InstanceNorm2d(out_filters))
+            layers.append(nn.LeakyReLU(0.2, inplace=True))
+            return layers
+
+        self.model = nn.Sequential(
+            *discriminator_block(channels, 64, normalize=False),
+            *discriminator_block(64, 128),
+            *discriminator_block(128, 256),
+            *discriminator_block(256, 512),
+            nn.ZeroPad2d((1, 0, 1, 0)),
+            nn.Conv2d(512, 1, 4, padding=1)
+        )
+
+    def forward(self, img):
+        return self.model(img)

huggan/pytorch/cyclegan/train.py

@@ -0,0 +1,354 @@
+import argparse
+import os
+import numpy as np
+import itertools
+from pathlib import Path
+import datetime
+import time
+import sys
+
+from PIL import Image
+
+from torchvision.transforms import Compose, Resize, ToTensor, Normalize, RandomCrop, RandomHorizontalFlip
+from torchvision.utils import save_image, make_grid
+
+from torch.utils.data import DataLoader
+
+from modeling_cyclegan import GeneratorResNet, Discriminator
+
+from utils import ReplayBuffer, LambdaLR
+
+from datasets import load_dataset
+
+from accelerate import Accelerator
+
+import torch.nn as nn
+import torch
+
+def parse_args(args=None):
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--epoch", type=int, default=0, help="epoch to start training from")
+    parser.add_argument("--num_epochs", type=int, default=200, help="number of epochs of training")
+    parser.add_argument("--dataset_name", type=str, default="huggan/facades", help="name of the dataset")
+    parser.add_argument("--batch_size", type=int, default=1, help="size of the batches")
+    parser.add_argument("--lr", type=float, default=0.0002, help="adam: learning rate")
+    parser.add_argument("--beta1", type=float, default=0.5, help="adam: decay of first order momentum of gradient")
+    parser.add_argument("--beta2", type=float, default=0.999, help="adam: decay of first order momentum of gradient")
+    parser.add_argument("--decay_epoch", type=int, default=100, help="epoch from which to start lr decay")
+    parser.add_argument("--num_workers", type=int, default=8, help="Number of CPU threads to use during batch generation")
+    parser.add_argument("--image_size", type=int, default=256, help="Size of images for training")
+    parser.add_argument("--channels", type=int, default=3, help="Number of image channels")
+    parser.add_argument("--sample_interval", type=int, default=100, help="interval between saving generator outputs")
+    parser.add_argument("--checkpoint_interval", type=int, default=-1, help="interval between saving model checkpoints")
+    parser.add_argument("--n_residual_blocks", type=int, default=9, help="number of residual blocks in generator")
+    parser.add_argument("--lambda_cyc", type=float, default=10.0, help="cycle loss weight")
+    parser.add_argument("--lambda_id", type=float, default=5.0, help="identity loss weight")
+    parser.add_argument("--fp16", action="store_true", help="If passed, will use FP16 training.")
+    parser.add_argument(
+        "--mixed_precision",
+        type=str,
+        default="no",
+        choices=["no", "fp16", "bf16"],
+        help="Whether to use mixed precision. Choose"
+        "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
+        "and an Nvidia Ampere GPU.",
+    )
+    parser.add_argument("--cpu", action="store_true", help="If passed, will train on the CPU.")
+    parser.add_argument(
+            "--push_to_hub",
+            action="store_true",
+            help="Whether to push the model to the HuggingFace hub after training.",
+            )
+    parser.add_argument(
+        "--pytorch_dump_folder_path",
+        required="--push_to_hub" in sys.argv,
+        type=Path,
+        help="Path to save the model. Will be created if it doesn't exist already.",
+    )
+    parser.add_argument(
+        "--model_name",
+        required="--push_to_hub" in sys.argv,
+        type=str,
+        help="Name of the model on the hub.",
+    )
+    parser.add_argument(
+        "--organization_name",
+        required=False,
+        default="huggan",
+        type=str,
+        help="Organization name to push to, in case args.push_to_hub is specified.",
+    )
+    return parser.parse_args(args=args)
+
+
+def weights_init_normal(m):
+        classname = m.__class__.__name__
+        if classname.find("Conv") != -1:
+            torch.nn.init.normal_(m.weight.data, 0.0, 0.02)
+            if hasattr(m, "bias") and m.bias is not None:
+                torch.nn.init.constant_(m.bias.data, 0.0)
+        elif classname.find("BatchNorm2d") != -1:
+            torch.nn.init.normal_(m.weight.data, 1.0, 0.02)
+            torch.nn.init.constant_(m.bias.data, 0.0)
+
+
+def training_function(config, args):
+    accelerator = Accelerator(fp16=args.fp16, cpu=args.cpu, mixed_precision=args.mixed_precision)
+    
+    # Create sample and checkpoint directories
+    os.makedirs("images/%s" % args.dataset_name, exist_ok=True)
+    os.makedirs("saved_models/%s" % args.dataset_name, exist_ok=True)
+
+    # Losses
+    criterion_GAN = torch.nn.MSELoss()
+    criterion_cycle = torch.nn.L1Loss()
+    criterion_identity = torch.nn.L1Loss()
+
+    input_shape = (args.channels, args.image_size, args.image_size)
+    # Calculate output shape of image discriminator (PatchGAN)
+    output_shape = (1, args.image_size // 2 ** 4, args.image_size // 2 ** 4)
+
+    # Initialize generator and discriminator
+    G_AB = GeneratorResNet(input_shape, args.n_residual_blocks)
+    G_BA = GeneratorResNet(input_shape, args.n_residual_blocks)
+    D_A = Discriminator(args.channels)
+    D_B = Discriminator(args.channels)
+
+    if args.epoch != 0:
+        # Load pretrained models
+        G_AB.load_state_dict(torch.load("saved_models/%s/G_AB_%d.pth" % (args.dataset_name, args.epoch)))
+        G_BA.load_state_dict(torch.load("saved_models/%s/G_BA_%d.pth" % (args.dataset_name, args.epoch)))
+        D_A.load_state_dict(torch.load("saved_models/%s/D_A_%d.pth" % (args.dataset_name, args.epoch)))
+        D_B.load_state_dict(torch.load("saved_models/%s/D_B_%d.pth" % (args.dataset_name, args.epoch)))
+    else:
+        # Initialize weights
+        G_AB.apply(weights_init_normal)
+        G_BA.apply(weights_init_normal)
+        D_A.apply(weights_init_normal)
+        D_B.apply(weights_init_normal)
+
+    # Optimizers
+    optimizer_G = torch.optim.Adam(
+        itertools.chain(G_AB.parameters(), G_BA.parameters()), lr=args.lr, betas=(args.beta1, args.beta2)
+    )
+    optimizer_D_A = torch.optim.Adam(D_A.parameters(), lr=args.lr, betas=(args.beta1, args.beta2))
+    optimizer_D_B = torch.optim.Adam(D_B.parameters(), lr=args.lr, betas=(args.beta1, args.beta2))
+
+    # Learning rate update schedulers
+    lr_scheduler_G = torch.optim.lr_scheduler.LambdaLR(
+        optimizer_G, lr_lambda=LambdaLR(args.num_epochs, args.epoch, args.decay_epoch).step
+    )
+    lr_scheduler_D_A = torch.optim.lr_scheduler.LambdaLR(
+        optimizer_D_A, lr_lambda=LambdaLR(args.num_epochs, args.epoch, args.decay_epoch).step
+    )
+    lr_scheduler_D_B = torch.optim.lr_scheduler.LambdaLR(
+        optimizer_D_B, lr_lambda=LambdaLR(args.num_epochs, args.epoch, args.decay_epoch).step
+    )
+
+    # Buffers of previously generated samples
+    fake_A_buffer = ReplayBuffer()
+    fake_B_buffer = ReplayBuffer()
+
+    # Image transformations
+    transform = Compose([
+        Resize(int(args.image_size * 1.12), Image.BICUBIC),
+        RandomCrop((args.image_size, args.image_size)),
+        RandomHorizontalFlip(),
+        ToTensor(),
+        Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
+    ])
+
+    def transforms(examples):
+        examples["A"] = [transform(image.convert("RGB")) for image in examples["imageA"]]
+        examples["B"] = [transform(image.convert("RGB")) for image in examples["imageB"]]
+
+        del examples["imageA"]
+        del examples["imageB"]
+
+        return examples
+
+    dataset = load_dataset(args.dataset_name)
+    transformed_dataset = dataset.with_transform(transforms)
+
+    splits = transformed_dataset['train'].train_test_split(test_size=0.1)
+    train_ds = splits['train']
+    val_ds = splits['test']
+
+    dataloader = DataLoader(train_ds, shuffle=True, batch_size=args.batch_size, num_workers=args.num_workers)
+    val_dataloader = DataLoader(val_ds, batch_size=5, shuffle=True, num_workers=1)
+
+    def sample_images(batches_done):
+        """Saves a generated sample from the test set"""
+        batch = next(iter(val_dataloader))
+        G_AB.eval()
+        G_BA.eval()
+        real_A = batch["A"]
+        fake_B = G_AB(real_A)
+        real_B = batch["B"]
+        fake_A = G_BA(real_B)
+        # Arange images along x-axis
+        real_A = make_grid(real_A, nrow=5, normalize=True)
+        real_B = make_grid(real_B, nrow=5, normalize=True)
+        fake_A = make_grid(fake_A, nrow=5, normalize=True)
+        fake_B = make_grid(fake_B, nrow=5, normalize=True)
+        # Arange images along y-axis
+        image_grid = torch.cat((real_A, fake_B, real_B, fake_A), 1)
+        save_image(image_grid, "images/%s/%s.png" % (args.dataset_name, batches_done), normalize=False)
+
+    G_AB, G_BA, D_A, D_B, optimizer_G, optimizer_D_A, optimizer_D_B, dataloader, val_dataloader = accelerator.prepare(G_AB, G_BA, D_A, D_B, optimizer_G, optimizer_D_A, optimizer_D_B, dataloader, val_dataloader)
+    
+    # ----------
+    #  Training
+    # ----------
+
+    prev_time = time.time()
+    for epoch in range(args.epoch, args.num_epochs):
+        for i, batch in enumerate(dataloader):
+
+            # Set model input
+            real_A = batch["A"]
+            real_B = batch["B"]
+
+            # Adversarial ground truths
+            valid = torch.ones((real_A.size(0), *output_shape), device=accelerator.device)
+            fake = torch.zeros((real_A.size(0), *output_shape), device=accelerator.device)
+
+            # ------------------
+            #  Train Generators
+            # ------------------
+
+            G_AB.train()
+            G_BA.train()
+
+            optimizer_G.zero_grad()
+
+            # Identity loss
+            loss_id_A = criterion_identity(G_BA(real_A), real_A)
+            loss_id_B = criterion_identity(G_AB(real_B), real_B)
+
+            loss_identity = (loss_id_A + loss_id_B) / 2
+
+            # GAN loss
+            fake_B = G_AB(real_A)
+            loss_GAN_AB = criterion_GAN(D_B(fake_B), valid)
+            fake_A = G_BA(real_B)
+            loss_GAN_BA = criterion_GAN(D_A(fake_A), valid)
+
+            loss_GAN = (loss_GAN_AB + loss_GAN_BA) / 2
+
+            # Cycle loss
+            recov_A = G_BA(fake_B)
+            loss_cycle_A = criterion_cycle(recov_A, real_A)
+            recov_B = G_AB(fake_A)
+            loss_cycle_B = criterion_cycle(recov_B, real_B)
+
+            loss_cycle = (loss_cycle_A + loss_cycle_B) / 2
+
+            # Total loss
+            loss_G = loss_GAN + args.lambda_cyc * loss_cycle + args.lambda_id * loss_identity
+
+            accelerator.backward(loss_G)
+            optimizer_G.step()
+
+            # -----------------------
+            #  Train Discriminator A
+            # -----------------------
+
+            optimizer_D_A.zero_grad()
+
+            # Real loss
+            loss_real = criterion_GAN(D_A(real_A), valid)
+            # Fake loss (on batch of previously generated samples)
+            fake_A_ = fake_A_buffer.push_and_pop(fake_A)
+            loss_fake = criterion_GAN(D_A(fake_A_.detach()), fake)
+            # Total loss
+            loss_D_A = (loss_real + loss_fake) / 2
+
+            accelerator.backward(loss_D_A)
+            optimizer_D_A.step()
+
+            # -----------------------
+            #  Train Discriminator B
+            # -----------------------
+
+            optimizer_D_B.zero_grad()
+
+            # Real loss
+            loss_real = criterion_GAN(D_B(real_B), valid)
+            # Fake loss (on batch of previously generated samples)
+            fake_B_ = fake_B_buffer.push_and_pop(fake_B)
+            loss_fake = criterion_GAN(D_B(fake_B_.detach()), fake)
+            # Total loss
+            loss_D_B = (loss_real + loss_fake) / 2
+
+            accelerator.backward(loss_D_B)
+            optimizer_D_B.step()
+
+            loss_D = (loss_D_A + loss_D_B) / 2
+
+            # --------------
+            #  Log Progress
+            # --------------
+
+            # Determine approximate time left
+            batches_done = epoch * len(dataloader) + i
+            batches_left = args.num_epochs * len(dataloader) - batches_done
+            time_left = datetime.timedelta(seconds=batches_left * (time.time() - prev_time))
+            prev_time = time.time()
+
+            # Print log
+            sys.stdout.write(
+                "\r[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f, adv: %f, cycle: %f, identity: %f] ETA: %s"
+                % (
+                    epoch,
+                    args.num_epochs,
+                    i,
+                    len(dataloader),
+                    loss_D.item(),
+                    loss_G.item(),
+                    loss_GAN.item(),
+                    loss_cycle.item(),
+                    loss_identity.item(),
+                    time_left,
+                )
+            )
+
+            # If at sample interval save image
+            if batches_done % args.sample_interval == 0:
+                sample_images(batches_done)
+
+        # Update learning rates
+        lr_scheduler_G.step()
+        lr_scheduler_D_A.step()
+        lr_scheduler_D_B.step()
+
+        if args.checkpoint_interval != -1 and epoch % args.checkpoint_interval == 0:
+            # Save model checkpoints
+            torch.save(G_AB.state_dict(), "saved_models/%s/G_AB_%d.pth" % (args.dataset_name, epoch))
+            torch.save(G_BA.state_dict(), "saved_models/%s/G_BA_%d.pth" % (args.dataset_name, epoch))
+            torch.save(D_A.state_dict(), "saved_models/%s/D_A_%d.pth" % (args.dataset_name, epoch))
+            torch.save(D_B.state_dict(), "saved_models/%s/D_B_%d.pth" % (args.dataset_name, epoch))
+
+    # Optionally push to hub
+    if args.push_to_hub:
+        save_directory = args.pytorch_dump_folder_path
+        if not save_directory.exists():
+            save_directory.mkdir(parents=True)
+
+        G_AB.push_to_hub(
+            repo_path_or_name=save_directory / args.model_name,
+            organization=args.organization_name,
+        )
+
+def main():
+    args = parse_args()
+    print(args)
+
+    # Make directory for saving generated images
+    os.makedirs("images", exist_ok=True)
+
+    training_function({}, args)
+
+
+if __name__ == "__main__":
+    main()

huggan/pytorch/cyclegan/utils.py

@@ -0,0 +1,44 @@
+import random
+import time
+import datetime
+import sys
+
+from torch.autograd import Variable
+import torch
+import numpy as np
+
+from torchvision.utils import save_image
+
+
+class ReplayBuffer:
+    def __init__(self, max_size=50):
+        assert max_size > 0, "Empty buffer or trying to create a black hole. Be careful."
+        self.max_size = max_size
+        self.data = []
+
+    def push_and_pop(self, data):
+        to_return = []
+        for element in data.data:
+            element = torch.unsqueeze(element, 0)
+            if len(self.data) < self.max_size:
+                self.data.append(element)
+                to_return.append(element)
+            else:
+                if random.uniform(0, 1) > 0.5:
+                    i = random.randint(0, self.max_size - 1)
+                    to_return.append(self.data[i].clone())
+                    self.data[i] = element
+                else:
+                    to_return.append(element)
+        return Variable(torch.cat(to_return))
+
+
+class LambdaLR:
+    def __init__(self, n_epochs, offset, decay_start_epoch):
+        assert (n_epochs - decay_start_epoch) > 0, "Decay must start before the training session ends!"
+        self.n_epochs = n_epochs
+        self.offset = offset
+        self.decay_start_epoch = decay_start_epoch
+
+    def step(self, epoch):
+        return 1.0 - max(0, epoch + self.offset - self.decay_start_epoch) / (self.n_epochs - self.decay_start_epoch)

huggan/pytorch/dcgan/README.md

@@ -0,0 +1,155 @@
+# Train DCGAN on your custom data
+
+This folder contains a script to train [DCGAN](https://arxiv.org/abs/1511.06434) for unconditional image generation, leveraging the [Hugging Face](https://huggingface.co/) ecosystem for processing your data and pushing the model to the Hub.
+
+The script leverages 🤗 Datasets for loading and processing data, and 🤗 Accelerate for instantly running on CPU, single, multi-GPUs or TPU, also supporting fp16/mixed precision.
+
+<p align="center">
+    <img src="https://raw.githubusercontent.com/huggingface/community-events/main/huggan/assets/dcgan_mnist.png" alt="drawing" width="300"/>
+</p>
+
+
+## Launching the script
+
+To train the model with the default parameters (5 epochs, 64x64 images, etc.) on [MNIST](https://huggingface.co/datasets/mnist), first run:
+
+```bash
+accelerate config
+```
+
+and answer the questions asked about your environment. Next, launch the script as follows: 
+
+```bash
+accelerate launch train.py
+```
+
+This will create a local "images" directory, containing generated images over the course of the training.
+
+To train on another dataset available on the hub, simply do (for instance):
+
+```bash
+python train.py --dataset cifar-10
+```
+
+In case you'd like to tweak the script to your liking, first fork the "community-events" [repo](https://github.com/huggingface/community-events) (see the button on the top right), then clone it locally:
+
+```bash
+git clone https://github.com/<your Github username>/community-events.git
+```
+
+and edit to your liking.
+
+## Training on your own data
+
+You can of course also train on your own images. For this, one can leverage Datasets' [ImageFolder](https://huggingface.co/docs/datasets/v2.0.0/en/image_process#imagefolder). Make sure to authenticate with the hub first, by running the `huggingface-cli login` command in a terminal, or the following in case you're working in a notebook:
+
+```python
+from huggingface_hub import notebook_login
+
+notebook_login()
+```
+
+Next, run the following in a notebook/script:
+
+```python
+from datasets import load_dataset
+
+# first: load dataset
+# option 1: from local folder
+dataset = load_dataset("imagefolder", data_dir="path_to_folder")
+# option 2: from remote URL (e.g. a zip file)
+dataset = load_dataset("imagefolder", data_files="URL to .zip file")
+
+# next: push to the hub (assuming git-LFS is installed)
+dataset.push_to_hub("huggan/my-awesome-dataset")
+```
+
+You can then simply pass the name of the dataset to the script:
+
+```bash
+accelerate launch train.py --dataset huggan/my-awesome-dataset
+```
+
+## Pushing model to the Hub
+
+You can push your trained generator to the hub after training by specifying the `push_to_hub` flag, along with a `model_name` and `pytorch_dump_folder_path`. 
+
+```bash
+accelerate launch train.py --push_to_hub --model_name dcgan-mnist
+```
+
+This is made possible by making the generator inherit from `PyTorchModelHubMixin`available in the `huggingface_hub` library. 
+
+This means that after training, generating a new image can be done as follows:
+
+```python
+import torch
+import torch.nn as nn
+from torchvision.transforms import ToPILImage
+from huggingface_hub import PyTorchModelHubMixin
+
+class Generator(nn.Module, PyTorchModelHubMixin):
+    def __init__(self, num_channels=3, latent_dim=100, hidden_size=64):
+        super(Generator, self).__init__()
+        self.model = nn.Sequential(
+            # input is Z, going into a convolution
+            nn.ConvTranspose2d(latent_dim, hidden_size * 8, 4, 1, 0, bias=False),
+            nn.BatchNorm2d(hidden_size * 8),
+            nn.ReLU(True),
+            # state size. (hidden_size*8) x 4 x 4
+            nn.ConvTranspose2d(hidden_size * 8, hidden_size * 4, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(hidden_size * 4),
+            nn.ReLU(True),
+            # state size. (hidden_size*4) x 8 x 8
+            nn.ConvTranspose2d(hidden_size * 4, hidden_size * 2, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(hidden_size * 2),
+            nn.ReLU(True),
+            # state size. (hidden_size*2) x 16 x 16
+            nn.ConvTranspose2d(hidden_size * 2, hidden_size, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(hidden_size),
+            nn.ReLU(True),
+            # state size. (hidden_size) x 32 x 32
+            nn.ConvTranspose2d(hidden_size, num_channels, 4, 2, 1, bias=False),
+            nn.Tanh()
+            # state size. (num_channels) x 64 x 64
+        )
+
+    def forward(self, noise):
+        pixel_values = self.model(noise)
+
+        return pixel_values
+
+model = Generator.from_pretrained("huggan/dcgan-mnist")
+
+device = "cuda" if torch.cuda.is_available() else "cpu
+model.to(device)
+ 
+with torch.no_grad():
+    z = torch.randn(1, 100, 1, 1, device=device)
+    pixel_values = model(z)
+
+# turn into actual image
+image = pixel_values[0]
+image = (image + 1) /2
+image = ToPILImage()(image)
+image.save("generated.png")
+```
+
+## Weights and Biases integration
+
+You can easily add logging to [Weights and Biases](https://wandb.ai/site) by passing the `--wandb` flag:
+
+```bash
+accelerate launch train.py --wandb
+````
+
+You can then follow the progress of your GAN in a browser:
+
+<p align="center">
+    <img src="https://raw.githubusercontent.com/huggingface/community-events/main/huggan/assets/wandb.png" alt="drawing" width="700"/>
+</p>
+
+
+# Citation
+
+This repo is entirely based on PyTorch's official [DCGAN tutorial](https://pytorch.org/tutorials/beginner/dcgan_faces_tutorial.html), but with added HuggingFace goodies.

huggan/pytorch/dcgan/modeling_dcgan.py

@@ -0,0 +1,80 @@
+#!/usr/bin/env python
+# coding=utf-8
+# Copyright (c) 2022 PyTorch contributors and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions.
+
+import torch.nn as nn
+
+from huggan.pytorch.huggan_mixin import HugGANModelHubMixin
+
+
+class Generator(nn.Module, HugGANModelHubMixin):
+    def __init__(self, num_channels=3, latent_dim=100, hidden_size=64):
+        super(Generator, self).__init__()
+        self.model = nn.Sequential(
+            # input is Z, going into a convolution
+            nn.ConvTranspose2d(latent_dim, hidden_size * 8, 4, 1, 0, bias=False),
+            nn.BatchNorm2d(hidden_size * 8),
+            nn.ReLU(True),
+            # state size. (hidden_size*8) x 4 x 4
+            nn.ConvTranspose2d(hidden_size * 8, hidden_size * 4, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(hidden_size * 4),
+            nn.ReLU(True),
+            # state size. (hidden_size*4) x 8 x 8
+            nn.ConvTranspose2d(hidden_size * 4, hidden_size * 2, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(hidden_size * 2),
+            nn.ReLU(True),
+            # state size. (hidden_size*2) x 16 x 16
+            nn.ConvTranspose2d(hidden_size * 2, hidden_size, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(hidden_size),
+            nn.ReLU(True),
+            # state size. (hidden_size) x 32 x 32
+            nn.ConvTranspose2d(hidden_size, num_channels, 4, 2, 1, bias=False),
+            nn.Tanh()
+            # state size. (num_channels) x 64 x 64
+        )
+
+    def forward(self, noise):
+        pixel_values = self.model(noise)
+
+        return pixel_values
+
+
+class Discriminator(nn.Module):
+    def __init__(self, num_channels=3, hidden_size=64):
+        super(Discriminator, self).__init__()
+        self.model = nn.Sequential(
+            # input is (num_channels) x 64 x 64
+            nn.Conv2d(num_channels, hidden_size, 4, 2, 1, bias=False),
+            nn.LeakyReLU(0.2, inplace=True),
+            # state size. (hidden_size) x 32 x 32
+            nn.Conv2d(hidden_size, hidden_size * 2, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(hidden_size * 2),
+            nn.LeakyReLU(0.2, inplace=True),
+            # state size. (hidden_size*2) x 16 x 16
+            nn.Conv2d(hidden_size * 2, hidden_size * 4, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(hidden_size * 4),
+            nn.LeakyReLU(0.2, inplace=True),
+            # state size. (hidden_size*4) x 8 x 8
+            nn.Conv2d(hidden_size * 4, hidden_size * 8, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(hidden_size * 8),
+            nn.LeakyReLU(0.2, inplace=True),
+            # state size. (hidden_size*8) x 4 x 4
+            nn.Conv2d(hidden_size * 8, 1, 4, 1, 0, bias=False),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, pixel_values):
+        logits = self.model(pixel_values)
+
+        return logits

huggan/pytorch/dcgan/train.py

@@ -0,0 +1,346 @@
+#!/usr/bin/env python
+# coding=utf-8
+# Copyright (c) 2022 PyTorch contributors and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions.
+
+""" Training a Deep Convolutional Generative Adversarial Network (DCGAN) leveraging the 🤗 ecosystem.
+Paper: https://arxiv.org/abs/1511.06434.
+Based on PyTorch's official tutorial: https://pytorch.org/tutorials/beginner/dcgan_faces_tutorial.html.
+"""
+
+
+import argparse
+import logging
+import os
+import sys
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+from torchvision.transforms import (CenterCrop, Compose, Normalize, Resize,
+                                    ToTensor, ToPILImage)
+from torchvision.utils import save_image
+
+from PIL import Image, ImageFile
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
+from accelerate import Accelerator
+
+from modeling_dcgan import Discriminator, Generator
+
+from datasets import load_dataset
+
+from huggan.pytorch.metrics.inception import InceptionV3
+from huggan.pytorch.metrics.fid_score import calculate_fretchet
+
+import wandb
+
+logger = logging.getLogger(__name__)
+
+
+def parse_args(args=None):
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dataset", type=str, default="mnist", help="Dataset to load from the HuggingFace hub.")
+    parser.add_argument("--num_workers", type=int, default=0, help="Number of workers when loading data")
+    parser.add_argument("--batch_size", type=int, default=128, help="Batch size to use during training")
+    parser.add_argument(
+        "--image_size",
+        type=int,
+        default=64,
+        help="Spatial size to use when resizing images for training.",
+    )
+    parser.add_argument(
+        "--num_channels",
+        type=int,
+        default=3,
+        help="Number of channels in the training images. For color images this is 3.",
+    )
+    parser.add_argument("--latent_dim", type=int, default=100, help="Dimensionality of the latent space.")
+    parser.add_argument(
+        "--generator_hidden_size",
+        type=int,
+        default=64,
+        help="Hidden size of the generator's feature maps.",
+    )
+    parser.add_argument(
+        "--discriminator_hidden_size",
+        type=int,
+        default=64,
+        help="Hidden size of the discriminator's feature maps.",
+    )
+    parser.add_argument("--num_epochs", type=int, default=5, help="number of epochs of training")
+    parser.add_argument("--lr", type=float, default=0.0002, help="adam: learning rate")
+    parser.add_argument(
+        "--beta1",
+        type=float,
+        default=0.5,
+        help="adam: decay of first order momentum of gradient",
+    )
+    parser.add_argument("--fp16", action="store_true", help="If passed, will use FP16 training.")
+    parser.add_argument(
+        "--mixed_precision",
+        type=str,
+        default="no",
+        choices=["no", "fp16", "bf16"],
+        help="Whether to use mixed precision. Choose"
+        "between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."
+        "and an Nvidia Ampere GPU.",
+    )
+    parser.add_argument("--cpu", action="store_true", help="If passed, will train on the CPU.")
+    parser.add_argument("--output_dir", type=Path, default=Path("./output"), help="Name of the directory to dump generated images during training.")
+    parser.add_argument("--wandb", action="store_true", help="If passed, will log to Weights and Biases.")
+    parser.add_argument(
+        "--logging_steps",
+        type=int,
+        default=50,
+        help="Number of steps between each logging",
+    )
+    parser.add_argument(
+        "--push_to_hub",
+        action="store_true",
+        help="Whether to push the model to the HuggingFace hub after training.",
+        )
+    parser.add_argument(
+        "--model_name",
+        default=None,
+        type=str,
+        help="Name of the model on the hub.",
+    )
+    parser.add_argument(
+        "--organization_name",
+        default="huggan",
+        type=str,
+        help="Organization name to push to, in case args.push_to_hub is specified.",
+    )
+    args = parser.parse_args()
+    
+    if args.push_to_hub:
+        assert args.output_dir is not None, "Need an `output_dir` to create a repo when `--push_to_hub` is passed."
+        assert args.model_name is not None, "Need a `model_name` to create a repo when `--push_to_hub` is passed."
+
+    if args.output_dir is not None:
+        os.makedirs(args.output_dir, exist_ok=True)
+    
+    return args
+
+
+# Custom weights initialization called on Generator and Discriminator
+def weights_init(m):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        nn.init.normal_(m.weight.data, 0.0, 0.02)
+    elif classname.find("BatchNorm") != -1:
+        nn.init.normal_(m.weight.data, 1.0, 0.02)
+        nn.init.constant_(m.bias.data, 0)
+
+
+def training_function(config, args):
+
+    # Initialize the accelerator. We will let the accelerator handle device placement for us in this example.
+    accelerator = Accelerator(fp16=args.fp16, cpu=args.cpu, mixed_precision=args.mixed_precision)
+
+    # Setup logging, we only want one process per machine to log things on the screen.
+    # accelerator.is_local_main_process is only True for one process per machine.
+    logger.setLevel(logging.INFO if accelerator.is_local_main_process else logging.ERROR)
+    if accelerator.is_local_main_process:
+        # set up Weights and Biases if requested
+        if args.wandb:
+            import wandb
+
+            wandb.init(project=str(args.output_dir).split("/")[-1])   
+    
+    # Loss function
+    criterion = nn.BCELoss()
+
+    # Initialize generator and discriminator
+    generator = Generator(
+        num_channels=args.num_channels,
+        latent_dim=args.latent_dim,
+        hidden_size=args.generator_hidden_size,
+    )
+    discriminator = Discriminator(num_channels=args.num_channels, hidden_size=args.discriminator_hidden_size)
+
+    # Initialize weights
+    generator.apply(weights_init)
+    discriminator.apply(weights_init)
+
+    # Initialize Inceptionv3 (for FID metric)
+    model = InceptionV3()
+
+    # Initialize Inceptionv3 (for FID metric)
+    model = InceptionV3()
+
+    # Create batch of latent vectors that we will use to visualize
+    # the progression of the generator
+    fixed_noise = torch.randn(64, args.latent_dim, 1, 1, device=accelerator.device)
+
+    # Establish convention for real and fake labels during training
+    real_label = 1.0
+    fake_label = 0.0
+
+    # Setup Adam optimizers for both G and D
+    discriminator_optimizer = torch.optim.Adam(discriminator.parameters(), lr=args.lr, betas=(args.beta1, 0.999))
+    generator_optimizer = torch.optim.Adam(generator.parameters(), lr=args.lr, betas=(args.beta1, 0.999))
+
+    # Configure data loader
+    dataset = load_dataset(args.dataset)
+
+    transform = Compose(
+        [
+            Resize(args.image_size),
+            CenterCrop(args.image_size),
+            ToTensor(),
+            Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
+        ]
+    )
+
+    def transforms(examples):
+        examples["pixel_values"] = [transform(image.convert("RGB")) for image in examples["image"]]
+
+        del examples["image"]
+
+        return examples
+
+    transformed_dataset = dataset.with_transform(transforms)
+
+    dataloader = DataLoader(
+        transformed_dataset["train"], batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers
+    )
+
+    generator, discriminator, generator_optimizer, discriminator_optimizer, dataloader = accelerator.prepare(generator, discriminator, generator_optimizer, discriminator_optimizer, dataloader)
+
+    # ----------
+    #  Training
+    # ----------
+    
+    # Training Loop
+
+    # Lists to keep track of progress
+    img_list = []
+
+    logger.info("***** Running training *****")
+    logger.info(f"  Num Epochs = {args.num_epochs}")
+    # For each epoch
+    for epoch in range(args.num_epochs):
+        # For each batch in the dataloader
+        for step, batch in enumerate(dataloader, 0):
+
+            ############################
+            # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))
+            ###########################
+            ## Train with all-real batch
+            discriminator.zero_grad()
+            # Format batch
+            real_cpu = batch["pixel_values"]
+            batch_size = real_cpu.size(0)
+            label = torch.full((batch_size,), real_label, dtype=torch.float, device=accelerator.device)
+            # Forward pass real batch through D
+            output = discriminator(real_cpu).view(-1)
+            # Calculate loss on all-real batch
+            errD_real = criterion(output, label)
+            # Calculate gradients for D in backward pass
+            accelerator.backward(errD_real)
+            D_x = output.mean().item()
+
+            ## Train with all-fake batch
+            # Generate batch of latent vectors
+            noise = torch.randn(batch_size, args.latent_dim, 1, 1, device=accelerator.device)
+            # Generate fake image batch with G
+            fake = generator(noise)
+            label.fill_(fake_label)
+            # Classify all fake batch with D
+            output = discriminator(fake.detach()).view(-1)
+            # Calculate D's loss on the all-fake batch
+            errD_fake = criterion(output, label)
+            # Calculate the gradients for this batch, accumulated (summed) with previous gradients
+            accelerator.backward(errD_fake)
+            D_G_z1 = output.mean().item()
+            # Compute error of D as sum over the fake and the real batches
+            errD = errD_real + errD_fake
+            # Update D
+            discriminator_optimizer.step()
+
+            ############################
+            # (2) Update G network: maximize log(D(G(z)))
+            ###########################
+            generator.zero_grad()
+            label.fill_(real_label)  # fake labels are real for generator cost
+            # Since we just updated D, perform another forward pass of all-fake batch through D
+            output = discriminator(fake).view(-1)
+            # Calculate G's loss based on this output
+            errG = criterion(output, label)
+            # Calculate gradients for G
+            accelerator.backward(errG)
+            D_G_z2 = output.mean().item()
+            # Update G
+            generator_optimizer.step()
+
+            # Log all results
+            if (step + 1) % args.logging_steps == 0:
+                errD.detach()
+                errG.detach()
+
+                if accelerator.state.num_processes > 1:
+                    errD = accelerator.gather(errD).sum() / accelerator.state.num_processes
+                    errG = accelerator.gather(errG).sum() / accelerator.state.num_processes
+
+                    train_logs = {
+                        "epoch": epoch,
+                        "discriminator_loss": errD,
+                        "generator_loss": errG,
+                        "D_x": D_x,
+                        "D_G_z1": D_G_z1,
+                        "D_G_z2": D_G_z2,
+                    }
+                    log_str = ""
+                    for k, v in train_logs.items():
+                        log_str += "| {}: {:.3e}".format(k, v)
+
+                    if accelerator.is_local_main_process:
+                        logger.info(log_str)
+                        if args.wandb:
+                            wandb.log(train_logs)
+
+            # Check how the generator is doing by saving G's output on fixed_noise
+            if (step % 500 == 0) or ((epoch == args.num_epochs - 1) and (step == len(dataloader) - 1)):
+                with torch.no_grad():
+                    fake_images = generator(fixed_noise).detach().cpu()
+                file_name = args.output_dir/f"iter_{step}.png"
+                save_image(fake_images.data[:25], file_name, nrow=5, normalize=True)
+                if accelerator.is_local_main_process and args.wandb:
+                    wandb.log({'generated_examples': wandb.Image(str(file_name)) })
+
+        # Calculate FID metric
+        fid = calculate_fretchet(real_cpu, fake, model.to(accelerator.device))
+        logger.info(f"FID: {fid}")
+        if accelerator.is_local_main_process and args.wandb:
+            wandb.log({"FID": fid})
+
+    # Optionally push to hub
+    if accelerator.is_main_process and args.push_to_hub:
+        generator.module.push_to_hub(
+            repo_path_or_name=args.output_dir / args.model_name,
+            organization=args.organization_name,
+        )
+
+
+def main():
+    args = parse_args()
+    print(args)
+
+    training_function({}, args)
+
+
+if __name__ == "__main__":
+    main()

huggan/pytorch/huggan_mixin.py

@@ -0,0 +1,131 @@
+from pathlib import Path
+from re import TEMPLATE
+from typing import Optional, Union
+import os
+
+from huggingface_hub import PyTorchModelHubMixin, HfApi, HfFolder, Repository
+
+from huggan import TEMPLATE_MODEL_CARD_PATH
+
+
+class HugGANModelHubMixin(PyTorchModelHubMixin):
+    """A mixin to push PyTorch Models to the Hugging Face Hub. This
+    mixin was adapted from the PyTorchModelHubMixin to also push a template
+    README.md for the HugGAN sprint.
+    """
+
+    def push_to_hub(
+        self,
+        repo_path_or_name: Optional[str] = None,
+        repo_url: Optional[str] = None,
+        commit_message: Optional[str] = "Add model",
+        organization: Optional[str] = None,
+        private: Optional[bool] = None,
+        api_endpoint: Optional[str] = None,
+        use_auth_token: Optional[Union[bool, str]] = None,
+        git_user: Optional[str] = None,
+        git_email: Optional[str] = None,
+        config: Optional[dict] = None,
+        skip_lfs_files: bool = False,
+        default_model_card: Optional[str] = TEMPLATE_MODEL_CARD_PATH
+    ) -> str:
+        """
+        Upload model checkpoint or tokenizer files to the Hub while
+        synchronizing a local clone of the repo in `repo_path_or_name`.
+        Parameters:
+            repo_path_or_name (`str`, *optional*):
+                Can either be a repository name for your model or tokenizer in
+                the Hub or a path to a local folder (in which case the
+                repository will have the name of that local folder). If not
+                specified, will default to the name given by `repo_url` and a
+                local directory with that name will be created.
+            repo_url (`str`, *optional*):
+                Specify this in case you want to push to an existing repository
+                in the hub. If unspecified, a new repository will be created in
+                your namespace (unless you specify an `organization`) with
+                `repo_name`.
+            commit_message (`str`, *optional*):
+                Message to commit while pushing. Will default to `"add config"`,
+                `"add tokenizer"` or `"add model"` depending on the type of the
+                class.
+            organization (`str`, *optional*):
+                Organization in which you want to push your model or tokenizer
+                (you must be a member of this organization).
+            private (`bool`, *optional*):
+                Whether the repository created should be private.
+            api_endpoint (`str`, *optional*):
+                The API endpoint to use when pushing the model to the hub.
+            use_auth_token (`bool` or `str`, *optional*):
+                The token to use as HTTP bearer authorization for remote files.
+                If `True`, will use the token generated when running
+                `transformers-cli login` (stored in `~/.huggingface`). Will
+                default to `True` if `repo_url` is not specified.
+            git_user (`str`, *optional*):
+                will override the `git config user.name` for committing and
+                pushing files to the hub.
+            git_email (`str`, *optional*):
+                will override the `git config user.email` for committing and
+                pushing files to the hub.
+            config (`dict`, *optional*):
+                Configuration object to be saved alongside the model weights.
+            default_model_card (`str`, *optional*):
+                Path to a markdown file to use as your default model card.
+        Returns:
+            The url of the commit of your model in the given repository.
+        """
+
+        if repo_path_or_name is None and repo_url is None:
+            raise ValueError(
+                "You need to specify a `repo_path_or_name` or a `repo_url`."
+            )
+
+        if use_auth_token is None and repo_url is None:
+            token = HfFolder.get_token()
+            if token is None:
+                raise ValueError(
+                    "You must login to the Hugging Face hub on this computer by typing `huggingface-cli login` and "
+                    "entering your credentials to use `use_auth_token=True`. Alternatively, you can pass your own "
+                    "token as the `use_auth_token` argument."
+                )
+        elif isinstance(use_auth_token, str):
+            token = use_auth_token
+        else:
+            token = None
+
+        if repo_path_or_name is None:
+            repo_path_or_name = repo_url.split("/")[-1]
+
+        # If no URL is passed and there's no path to a directory containing files, create a repo
+        if repo_url is None and not os.path.exists(repo_path_or_name):
+            repo_id = Path(repo_path_or_name).name
+            if organization:
+                repo_id = f"{organization}/{repo_id}"
+            repo_url = HfApi(endpoint=api_endpoint).create_repo(
+                repo_id=repo_id,
+                token=token,
+                private=private,
+                repo_type=None,
+                exist_ok=True,
+            )
+
+        repo = Repository(
+            repo_path_or_name,
+            clone_from=repo_url,
+            use_auth_token=use_auth_token,
+            git_user=git_user,
+            git_email=git_email,
+            skip_lfs_files=skip_lfs_files
+        )
+        repo.git_pull(rebase=True)
+
+        # Save the files in the cloned repo
+        self.save_pretrained(repo_path_or_name, config=config)
+
+        model_card_path = Path(repo_path_or_name) / 'README.md'
+        if not model_card_path.exists():
+            model_card_path.write_text(TEMPLATE_MODEL_CARD_PATH.read_text())
+
+        # Commit and push!
+        repo.git_add()
+        repo.git_commit(commit_message)
+        return repo.git_push()

huggan/pytorch/lightweight_gan/README.md

@@ -0,0 +1,89 @@
+# Train Lightweight GAN on your custom data
+
+This folder contains a script to train ['Lightweight' GAN](https://openreview.net/forum?id=1Fqg133qRaI) for unconditional image generation, leveraging the [Hugging Face](https://huggingface.co/) ecosystem for processing your data and pushing the model to the Hub.
+
+The script leverages 🤗 Datasets for loading and processing data, and 🤗 Accelerate for instantly running on CPU, single, multi-GPUs or TPU, also supporting mixed precision.
+
+<p align="center">
+    <img src="https://raw.githubusercontent.com/lucidrains/lightweight-gan/main/images/pizza-512.jpg" alt="drawing" width="300"/>
+</p>
+
+Pizza's that don't exist. Courtesy of Phil Wang.
+
+## Launching the script
+
+To train the model with the default parameters on [huggan/CelebA-faces](https://huggingface.co/datasets/huggan/CelebA-faces), first run:
+
+```bash
+accelerate config
+```
+
+and answer the questions asked about your environment. Next, launch the script as follows: 
+
+```bash
+accelerate launch cli.py
+```
+
+This will instantly run on multi-GPUs (if you asked for that). To train on another dataset available on the hub, simply do (for instance):
+
+```bash
+accelerate launch cli.py --dataset_name huggan/pokemon
+```
+
+In case you'd like to tweak the script to your liking, first fork the "community-events" [repo](https://github.com/huggingface/community-events) (see the button on the top right), then clone it locally:
+
+```bash
+git clone https://github.com/<your Github username>/community-events.git
+```
+
+and edit to your liking.
+
+## Training on your own data
+
+You can of course also train on your own images. For this, one can leverage Datasets' [ImageFolder](https://huggingface.co/docs/datasets/v2.0.0/en/image_process#imagefolder). Make sure to authenticate with the hub first, by running the `huggingface-cli login` command in a terminal, or the following in case you're working in a notebook:
+
+```python
+from huggingface_hub import notebook_login
+
+notebook_login()
+```
+
+Next, run the following in a notebook/script:
+
+```python
+from datasets import load_dataset
+
+# first: load dataset
+# option 1: from local folder
+dataset = load_dataset("imagefolder", data_dir="path_to_folder")
+# option 2: from remote URL (e.g. a zip file)
+dataset = load_dataset("imagefolder", data_files="URL to .zip file")
+
+# next: push to the hub (assuming git-LFS is installed)
+dataset.push_to_hub("huggan/my-awesome-dataset")
+```
+
+You can then simply pass the name of the dataset to the script:
+
+```bash
+accelerate launch cli.py --dataset huggan/my-awesome-dataset
+```
+
+## Weights and Biases integration
+
+You can easily add logging to [Weights and Biases](https://wandb.ai/site) by passing the `--wandb` flag:
+
+```bash
+accelerate launch cli.py --wandb
+````
+
+You can then follow the progress of your GAN in a browser:
+
+<p align="center">
+    <img src="https://raw.githubusercontent.com/huggingface/community-events/main/huggan/assets/lightweight_gan_wandb.png" alt="drawing" width="700"/>
+</p>
+
+
+# Citation
+
+This repo is entirely based on lucidrains' [Pytorch implementation](https://github.com/lucidrains/lightweight-gan), but with added HuggingFace goodies.

huggan/pytorch/lightweight_gan/cli.py

@@ -0,0 +1,178 @@
+import fire
+import random
+from retry.api import retry_call
+from tqdm import tqdm
+from datetime import datetime
+from pathlib import Path
+from lightweight_gan import Trainer, NanException
+
+import torch
+import torch.multiprocessing as mp
+
+import numpy as np
+
+def exists(val):
+    return val is not None
+
+def default(val, d):
+    return val if exists(val) else d
+
+def cast_list(el):
+    return el if isinstance(el, list) else [el]
+
+def timestamped_filename(prefix = 'generated-'):
+    now = datetime.now()
+    timestamp = now.strftime("%m-%d-%Y_%H-%M-%S")
+    return f'{prefix}{timestamp}'
+
+def set_seed(seed):
+    torch.manual_seed(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+    np.random.seed(seed)
+    random.seed(seed)
+
+def run_training(model_args, data, load_from, new, num_train_steps, name, seed):
+    
+    if seed is not None:
+        set_seed(seed)
+
+    model = Trainer(**model_args)
+
+    if not new:
+        model.load(load_from)
+    else:
+        model.clear()
+
+    progress_bar = tqdm(initial = model.steps, total = num_train_steps, mininterval=10., desc=f'{name}<{data}>')
+    G, D, D_aug = model.init_accelerator()
+
+    # model.set_data_src(data)
+
+    while model.steps < num_train_steps:
+        # retry_call(model.train, tries=3, exceptions=NanException)
+        model.train(G, D, D_aug)
+        progress_bar.n = model.steps
+        progress_bar.refresh()
+        if model.accelerator.is_local_main_process and model.steps % 50 == 0:
+            model.print_log()
+
+    model.save(model.checkpoint_num)
+
+def train_from_folder(
+    dataset_name = 'huggan/CelebA-faces',
+    data = './data',
+    results_dir = './results',
+    models_dir = './models',
+    name = 'default',
+    new = False,
+    load_from = -1,
+    image_size = 256,
+    optimizer = 'adam',
+    fmap_max = 512,
+    transparent = False,
+    greyscale = False,
+    batch_size = 10,
+    gradient_accumulate_every = 4,
+    num_train_steps = 150000,
+    learning_rate = 2e-4,
+    save_every = 10000,
+    evaluate_every = 1000,
+    generate = False,
+    generate_types = ['default', 'ema'],
+    generate_interpolation = False,
+    aug_test = False,
+    aug_prob=None,
+    aug_types=['cutout', 'translation'],
+    dataset_aug_prob=0.,
+    attn_res_layers = [32],
+    freq_chan_attn = False,
+    disc_output_size = 1,
+    dual_contrast_loss = False,
+    antialias = False,
+    interpolation_num_steps = 100,
+    save_frames = False,
+    num_image_tiles = None,
+    calculate_fid_every = None,
+    calculate_fid_num_images = 12800,
+    clear_fid_cache = False,
+    seed = 42,
+    cpu = False,
+    mixed_precision = "no",
+    show_progress = False,
+    wandb = False,
+    push_to_hub = False,
+    organization_name = None,
+):
+    if push_to_hub:
+        if name == 'default':
+            raise RuntimeError(
+                "You've chosen to push to hub, but have left the --name flag as 'default'."
+                " You should name your model something other than 'default'!"
+            )
+
+    num_image_tiles = default(num_image_tiles, 4 if image_size > 512 else 8)
+
+    model_args = dict(
+        dataset_name = dataset_name,
+        name = name,
+        results_dir = results_dir,
+        models_dir = models_dir,
+        batch_size = batch_size,
+        gradient_accumulate_every = gradient_accumulate_every,
+        attn_res_layers = cast_list(attn_res_layers),
+        freq_chan_attn = freq_chan_attn,
+        disc_output_size = disc_output_size,
+        dual_contrast_loss = dual_contrast_loss,
+        antialias = antialias,
+        image_size = image_size,
+        num_image_tiles = num_image_tiles,
+        optimizer = optimizer,
+        fmap_max = fmap_max,
+        transparent = transparent,
+        greyscale = greyscale,
+        lr = learning_rate,
+        save_every = save_every,
+        evaluate_every = evaluate_every,
+        aug_prob = aug_prob,
+        aug_types = cast_list(aug_types),
+        dataset_aug_prob = dataset_aug_prob,
+        calculate_fid_every = calculate_fid_every,
+        calculate_fid_num_images = calculate_fid_num_images,
+        clear_fid_cache = clear_fid_cache,
+        cpu = cpu,
+        mixed_precision = mixed_precision,
+        wandb = wandb,
+        push_to_hub = push_to_hub,
+        organization_name = organization_name
+    )
+
+    if generate:
+        model = Trainer(**model_args)
+        model.load(load_from)
+        samples_name = timestamped_filename()
+        checkpoint = model.checkpoint_num
+        dir_result = model.generate(samples_name, num_image_tiles, checkpoint, generate_types)
+        print(f'sample images generated at {dir_result}')
+        return
+
+    if generate_interpolation:
+        model = Trainer(**model_args)
+        model.load(load_from)
+        samples_name = timestamped_filename()
+        model.generate_interpolation(samples_name, num_image_tiles, num_steps = interpolation_num_steps, save_frames = save_frames)
+        print(f'interpolation generated at {results_dir}/{name}/{samples_name}')
+        return
+
+    if show_progress:
+        model = Trainer(**model_args)
+        model.show_progress(num_images=num_image_tiles, types=generate_types)
+        return
+
+    run_training(model_args, data, load_from, new, num_train_steps, name, seed)
+
+def main():
+    fire.Fire(train_from_folder)
+
+if __name__ == "__main__":
+    main()

huggan/pytorch/lightweight_gan/diff_augment.py

@@ -0,0 +1,102 @@
+import random
+
+import torch
+import torch.nn.functional as F
+
+
+def DiffAugment(x, types=[]):
+    for p in types:
+        for f in AUGMENT_FNS[p]:
+            x = f(x)
+    return x.contiguous()
+
+
+# """
+# Augmentation functions got images as `x`
+# where `x` is tensor with this dimensions:
+# 0 - count of images
+# 1 - channels
+# 2 - width
+# 3 - height of image
+# """
+
+def rand_brightness(x):
+    x = x + (torch.rand(x.size(0), 1, 1, 1, dtype=x.dtype, device=x.device) - 0.5)
+    return x
+
+def rand_saturation(x):
+    x_mean = x.mean(dim=1, keepdim=True)
+    x = (x - x_mean) * (torch.rand(x.size(0), 1, 1, 1, dtype=x.dtype, device=x.device) * 2) + x_mean
+    return x
+
+def rand_contrast(x):
+    x_mean = x.mean(dim=[1, 2, 3], keepdim=True)
+    x = (x - x_mean) * (torch.rand(x.size(0), 1, 1, 1, dtype=x.dtype, device=x.device) + 0.5) + x_mean
+    return x
+
+def rand_translation(x, ratio=0.125):
+    shift_x, shift_y = int(x.size(2) * ratio + 0.5), int(x.size(3) * ratio + 0.5)
+    translation_x = torch.randint(-shift_x, shift_x + 1, size=[x.size(0), 1, 1], device=x.device)
+    translation_y = torch.randint(-shift_y, shift_y + 1, size=[x.size(0), 1, 1], device=x.device)
+    grid_batch, grid_x, grid_y = torch.meshgrid(
+        torch.arange(x.size(0), dtype=torch.long, device=x.device),
+        torch.arange(x.size(2), dtype=torch.long, device=x.device),
+        torch.arange(x.size(3), dtype=torch.long, device=x.device),
+    indexing = 'ij')
+    grid_x = torch.clamp(grid_x + translation_x + 1, 0, x.size(2) + 1)
+    grid_y = torch.clamp(grid_y + translation_y + 1, 0, x.size(3) + 1)
+    x_pad = F.pad(x, [1, 1, 1, 1, 0, 0, 0, 0])
+    x = x_pad.permute(0, 2, 3, 1).contiguous()[grid_batch, grid_x, grid_y].permute(0, 3, 1, 2)
+    return x
+
+def rand_offset(x, ratio=1, ratio_h=1, ratio_v=1):
+    w, h = x.size(2), x.size(3)
+
+    imgs = []
+    for img in x.unbind(dim = 0):
+        max_h = int(w * ratio * ratio_h)
+        max_v = int(h * ratio * ratio_v)
+
+        value_h = random.randint(0, max_h) * 2 - max_h
+        value_v = random.randint(0, max_v) * 2 - max_v
+
+        if abs(value_h) > 0:
+            img = torch.roll(img, value_h, 2)
+
+        if abs(value_v) > 0:
+            img = torch.roll(img, value_v, 1)
+
+        imgs.append(img)
+
+    return torch.stack(imgs)
+
+def rand_offset_h(x, ratio=1):
+    return rand_offset(x, ratio=1, ratio_h=ratio, ratio_v=0)
+
+def rand_offset_v(x, ratio=1):
+    return rand_offset(x, ratio=1, ratio_h=0, ratio_v=ratio)
+
+def rand_cutout(x, ratio=0.5):
+    cutout_size = int(x.size(2) * ratio + 0.5), int(x.size(3) * ratio + 0.5)
+    offset_x = torch.randint(0, x.size(2) + (1 - cutout_size[0] % 2), size=[x.size(0), 1, 1], device=x.device)
+    offset_y = torch.randint(0, x.size(3) + (1 - cutout_size[1] % 2), size=[x.size(0), 1, 1], device=x.device)
+    grid_batch, grid_x, grid_y = torch.meshgrid(
+        torch.arange(x.size(0), dtype=torch.long, device=x.device),
+        torch.arange(cutout_size[0], dtype=torch.long, device=x.device),
+        torch.arange(cutout_size[1], dtype=torch.long, device=x.device),
+    indexing = 'ij')
+    grid_x = torch.clamp(grid_x + offset_x - cutout_size[0] // 2, min=0, max=x.size(2) - 1)
+    grid_y = torch.clamp(grid_y + offset_y - cutout_size[1] // 2, min=0, max=x.size(3) - 1)
+    mask = torch.ones(x.size(0), x.size(2), x.size(3), dtype=x.dtype, device=x.device)
+    mask[grid_batch, grid_x, grid_y] = 0
+    x = x * mask.unsqueeze(1)
+    return x
+
+AUGMENT_FNS = {
+    'color': [rand_brightness, rand_saturation, rand_contrast],
+    'offset': [rand_offset],
+    'offset_h': [rand_offset_h],
+    'offset_v': [rand_offset_v],
+    'translation': [rand_translation],
+    'cutout': [rand_cutout],
+}

huggan/pytorch/lightweight_gan/lightweight_gan.py

@@ -0,0 +1,1598 @@
+import os
+import json
+import tempfile
+from random import random
+import math
+from math import log2, floor
+from pathlib import Path
+from functools import partial
+from contextlib import contextmanager, ExitStack
+from pathlib import Path
+from shutil import rmtree
+
+import torch
+from torch.optim import Adam
+from torch import nn, einsum
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+from torch.autograd import grad as torch_grad
+
+from PIL import Image
+import torchvision
+from torchvision import transforms
+from torchvision.utils import save_image
+from kornia.filters import filter2d
+
+from huggan.pytorch.lightweight_gan.diff_augment import DiffAugment
+
+from tqdm import tqdm
+from einops import rearrange, reduce, repeat
+
+from datasets import load_dataset
+
+from accelerate import Accelerator, DistributedDataParallelKwargs
+from huggingface_hub import hf_hub_download, create_repo
+
+from huggan.pytorch.huggan_mixin import HugGANModelHubMixin
+from huggan.utils.hub import get_full_repo_name
+
+# constants
+
+# NUM_CORES = multiprocessing.cpu_count()
+EXTS = ['jpg', 'jpeg', 'png']
+PYTORCH_WEIGHTS_NAME = 'model.pt'
+
+
+# helpers
+
+def exists(val):
+    return val is not None
+
+
+@contextmanager
+def null_context():
+    yield
+
+
+def is_power_of_two(val):
+    return log2(val).is_integer()
+
+
+def default(val, d):
+    return val if exists(val) else d
+
+
+def set_requires_grad(model, bool):
+    for p in model.parameters():
+        p.requires_grad = bool
+
+
+def cycle(iterable):
+    while True:
+        for i in iterable:
+            yield i
+
+
+def raise_if_nan(t):
+    if torch.isnan(t):
+        raise NanException
+
+
+def evaluate_in_chunks(max_batch_size, model, *args):
+    split_args = list(zip(*list(map(lambda x: x.split(max_batch_size, dim=0), args))))
+    chunked_outputs = [model(*i) for i in split_args]
+    if len(chunked_outputs) == 1:
+        return chunked_outputs[0]
+    return torch.cat(chunked_outputs, dim=0)
+
+
+def slerp(val, low, high):
+    low_norm = low / torch.norm(low, dim=1, keepdim=True)
+    high_norm = high / torch.norm(high, dim=1, keepdim=True)
+    omega = torch.acos((low_norm * high_norm).sum(1))
+    so = torch.sin(omega)
+    res = (torch.sin((1.0 - val) * omega) / so).unsqueeze(1) * low + (torch.sin(val * omega) / so).unsqueeze(1) * high
+    return res
+
+
+def safe_div(n, d):
+    try:
+        res = n / d
+    except ZeroDivisionError:
+        prefix = '' if int(n >= 0) else '-'
+        res = float(f'{prefix}inf')
+    return res
+
+
+# loss functions
+
+def gen_hinge_loss(fake, real):
+    return fake.mean()
+
+
+def hinge_loss(real, fake):
+    return (F.relu(1 + real) + F.relu(1 - fake)).mean()
+
+
+def dual_contrastive_loss(real_logits, fake_logits):
+    device = real_logits.device
+    real_logits, fake_logits = map(lambda t: rearrange(t, '... -> (...)'), (real_logits, fake_logits))
+
+    def loss_half(t1, t2):
+        t1 = rearrange(t1, 'i -> i ()')
+        t2 = repeat(t2, 'j -> i j', i=t1.shape[0])
+        t = torch.cat((t1, t2), dim=-1)
+        return F.cross_entropy(t, torch.zeros(t1.shape[0], device=device, dtype=torch.long))
+
+    return loss_half(real_logits, fake_logits) + loss_half(-fake_logits, -real_logits)
+
+
+# helper classes
+
+class NanException(Exception):
+    pass
+
+
+class EMA():
+    def __init__(self, beta):
+        super().__init__()
+        self.beta = beta
+
+    def update_average(self, old, new):
+        if not exists(old):
+            return new
+        return old * self.beta + (1 - self.beta) * new
+
+
+class RandomApply(nn.Module):
+    def __init__(self, prob, fn, fn_else=lambda x: x):
+        super().__init__()
+        self.fn = fn
+        self.fn_else = fn_else
+        self.prob = prob
+
+    def forward(self, x):
+        fn = self.fn if random() < self.prob else self.fn_else
+        return fn(x)
+
+
+class ChanNorm(nn.Module):
+    def __init__(self, dim, eps=1e-5):
+        super().__init__()
+        self.eps = eps
+        self.g = nn.Parameter(torch.ones(1, dim, 1, 1))
+        self.b = nn.Parameter(torch.zeros(1, dim, 1, 1))
+
+    def forward(self, x):
+        var = torch.var(x, dim=1, unbiased=False, keepdim=True)
+        mean = torch.mean(x, dim=1, keepdim=True)
+        return (x - mean) / (var + self.eps).sqrt() * self.g + self.b
+
+
+class PreNorm(nn.Module):
+    def __init__(self, dim, fn):
+        super().__init__()
+        self.fn = fn
+        self.norm = ChanNorm(dim)
+
+    def forward(self, x):
+        return self.fn(self.norm(x))
+
+
+class Residual(nn.Module):
+    def __init__(self, fn):
+        super().__init__()
+        self.fn = fn
+
+    def forward(self, x):
+        return self.fn(x) + x
+
+
+class SumBranches(nn.Module):
+    def __init__(self, branches):
+        super().__init__()
+        self.branches = nn.ModuleList(branches)
+
+    def forward(self, x):
+        return sum(map(lambda fn: fn(x), self.branches))
+
+
+class Fuzziness(nn.Module):
+    def __init__(self):
+        super().__init__()
+        f = torch.Tensor([1, 2, 1])
+        self.register_buffer('f', f)
+
+    def forward(self, x):
+        f = self.f
+        f = f[None, None, :] * f[None, :, None]
+        return filter2d(x, f, normalized=True)
+
+
+Blur = nn.Identity
+
+
+# attention
+
+class DepthWiseConv2d(nn.Module):
+    def __init__(self, dim_in, dim_out, kernel_size, padding=0, stride=1, bias=True):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Conv2d(dim_in, dim_in, kernel_size=kernel_size, padding=padding, groups=dim_in, stride=stride,
+                      bias=bias),
+            nn.Conv2d(dim_in, dim_out, kernel_size=1, bias=bias)
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class LinearAttention(nn.Module):
+    def __init__(self, dim, dim_head=64, heads=8):
+        super().__init__()
+        self.scale = dim_head ** -0.5
+        self.heads = heads
+        inner_dim = dim_head * heads
+
+        self.nonlin = nn.GELU()
+        self.to_q = nn.Conv2d(dim, inner_dim, 1, bias=False)
+        self.to_kv = DepthWiseConv2d(dim, inner_dim * 2, 3, padding=1, bias=False)
+        self.to_out = nn.Conv2d(inner_dim, dim, 1)
+
+    def forward(self, fmap):
+        h, x, y = self.heads, *fmap.shape[-2:]
+        q, k, v = (self.to_q(fmap), *self.to_kv(fmap).chunk(2, dim=1))
+        q, k, v = map(lambda t: rearrange(t, 'b (h c) x y -> (b h) (x y) c', h=h), (q, k, v))
+
+        q = q.softmax(dim=-1)
+        k = k.softmax(dim=-2)
+
+        q = q * self.scale
+
+        context = einsum('b n d, b n e -> b d e', k, v)
+        out = einsum('b n d, b d e -> b n e', q, context)
+        out = rearrange(out, '(b h) (x y) d -> b (h d) x y', h=h, x=x, y=y)
+
+        out = self.nonlin(out)
+        return self.to_out(out)
+
+
+# dataset
+
+def convert_image_to(img_type, image):
+    if image.mode != img_type:
+        return image.convert(img_type)
+    return image
+
+
+class identity(object):
+    def __call__(self, tensor):
+        return tensor
+
+
+class expand_greyscale(object):
+    def __init__(self, transparent):
+        self.transparent = transparent
+
+    def __call__(self, tensor):
+        channels = tensor.shape[0]
+        num_target_channels = 4 if self.transparent else 3
+
+        if channels == num_target_channels:
+            return tensor
+
+        alpha = None
+        if channels == 1:
+            color = tensor.expand(3, -1, -1)
+        elif channels == 2:
+            color = tensor[:1].expand(3, -1, -1)
+            alpha = tensor[1:]
+        else:
+            raise Exception(f'image with invalid number of channels given {channels}')
+
+        if not exists(alpha) and self.transparent:
+            alpha = torch.ones(1, *tensor.shape[1:], device=tensor.device)
+
+        return color if not self.transparent else torch.cat((color, alpha))
+
+
+def resize_to_minimum_size(min_size, image):
+    if max(*image.size) < min_size:
+        return torchvision.transforms.functional.resize(image, min_size)
+    return image
+
+
+# augmentations
+
+def random_hflip(tensor, prob):
+    if prob > random():
+        return tensor
+    return torch.flip(tensor, dims=(3,))
+
+
+class AugWrapper(nn.Module):
+    def __init__(self, D, image_size):
+        super().__init__()
+        self.D = D
+
+    def forward(self, images, prob=0., types=[], detach=False, **kwargs):
+        context = torch.no_grad if detach else null_context
+
+        with context():
+            if random() < prob:
+                images = random_hflip(images, prob=0.5)
+                images = DiffAugment(images, types=types)
+
+        return self.D(images, **kwargs)
+
+
+# modifiable global variables
+
+norm_class = nn.BatchNorm2d
+
+
+def upsample(scale_factor=2):
+    return nn.Upsample(scale_factor=scale_factor)
+
+
+# squeeze excitation classes
+
+# global context network
+# https://arxiv.org/abs/2012.13375
+# similar to squeeze-excite, but with a simplified attention pooling and a subsequent layer norm
+
+class GlobalContext(nn.Module):
+    def __init__(
+            self,
+            *,
+            chan_in,
+            chan_out
+    ):
+        super().__init__()
+        self.to_k = nn.Conv2d(chan_in, 1, 1)
+        chan_intermediate = max(3, chan_out // 2)
+
+        self.net = nn.Sequential(
+            nn.Conv2d(chan_in, chan_intermediate, 1),
+            nn.LeakyReLU(0.1),
+            nn.Conv2d(chan_intermediate, chan_out, 1),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        context = self.to_k(x)
+        context = context.flatten(2).softmax(dim=-1)
+        out = einsum('b i n, b c n -> b c i', context, x.flatten(2))
+        out = out.unsqueeze(-1)
+        return self.net(out)
+
+
+# frequency channel attention
+# https://arxiv.org/abs/2012.11879
+
+def get_1d_dct(i, freq, L):
+    result = math.cos(math.pi * freq * (i + 0.5) / L) / math.sqrt(L)
+    return result * (1 if freq == 0 else math.sqrt(2))
+
+
+def get_dct_weights(width, channel, fidx_u, fidx_v):
+    dct_weights = torch.zeros(1, channel, width, width)
+    c_part = channel // len(fidx_u)
+
+    for i, (u_x, v_y) in enumerate(zip(fidx_u, fidx_v)):
+        for x in range(width):
+            for y in range(width):
+                coor_value = get_1d_dct(x, u_x, width) * get_1d_dct(y, v_y, width)
+                dct_weights[:, i * c_part: (i + 1) * c_part, x, y] = coor_value
+
+    return dct_weights
+
+
+class FCANet(nn.Module):
+    def __init__(
+            self,
+            *,
+            chan_in,
+            chan_out,
+            reduction=4,
+            width
+    ):
+        super().__init__()
+
+        freq_w, freq_h = ([0] * 8), list(range(8))  # in paper, it seems 16 frequencies was ideal
+        dct_weights = get_dct_weights(width, chan_in, [*freq_w, *freq_h], [*freq_h, *freq_w])
+        self.register_buffer('dct_weights', dct_weights)
+
+        chan_intermediate = max(3, chan_out // reduction)
+
+        self.net = nn.Sequential(
+            nn.Conv2d(chan_in, chan_intermediate, 1),
+            nn.LeakyReLU(0.1),
+            nn.Conv2d(chan_intermediate, chan_out, 1),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        x = reduce(x * self.dct_weights, 'b c (h h1) (w w1) -> b c h1 w1', 'sum', h1=1, w1=1)
+        return self.net(x)
+
+
+# generative adversarial network
+
+class Generator(nn.Module):
+    def __init__(
+            self,
+            *,
+            image_size,
+            latent_dim=256,
+            fmap_max=512,
+            fmap_inverse_coef=12,
+            transparent=False,
+            greyscale=False,
+            attn_res_layers=[],
+            freq_chan_attn=False
+    ):
+        super().__init__()
+        resolution = log2(image_size)
+        assert is_power_of_two(image_size), 'image size must be a power of 2'
+
+        if transparent:
+            init_channel = 4
+        elif greyscale:
+            init_channel = 1
+        else:
+            init_channel = 3
+
+        fmap_max = default(fmap_max, latent_dim)
+
+        self.initial_conv = nn.Sequential(
+            nn.ConvTranspose2d(latent_dim, latent_dim * 2, 4),
+            norm_class(latent_dim * 2),
+            nn.GLU(dim=1)
+        )
+
+        num_layers = int(resolution) - 2
+        features = list(map(lambda n: (n, 2 ** (fmap_inverse_coef - n)), range(2, num_layers + 2)))
+        features = list(map(lambda n: (n[0], min(n[1], fmap_max)), features))
+        features = list(map(lambda n: 3 if n[0] >= 8 else n[1], features))
+        features = [latent_dim, *features]
+
+        in_out_features = list(zip(features[:-1], features[1:]))
+
+        self.res_layers = range(2, num_layers + 2)
+        self.layers = nn.ModuleList([])
+        self.res_to_feature_map = dict(zip(self.res_layers, in_out_features))
+
+        self.sle_map = ((3, 7), (4, 8), (5, 9), (6, 10))
+        self.sle_map = list(filter(lambda t: t[0] <= resolution and t[1] <= resolution, self.sle_map))
+        self.sle_map = dict(self.sle_map)
+
+        self.num_layers_spatial_res = 1
+
+        for (res, (chan_in, chan_out)) in zip(self.res_layers, in_out_features):
+            image_width = 2 ** res
+
+            attn = None
+            if image_width in attn_res_layers:
+                attn = PreNorm(chan_in, LinearAttention(chan_in))
+
+            sle = None
+            if res in self.sle_map:
+                residual_layer = self.sle_map[res]
+                sle_chan_out = self.res_to_feature_map[residual_layer - 1][-1]
+
+                if freq_chan_attn:
+                    sle = FCANet(
+                        chan_in=chan_out,
+                        chan_out=sle_chan_out,
+                        width=2 ** (res + 1)
+                    )
+                else:
+                    sle = GlobalContext(
+                        chan_in=chan_out,
+                        chan_out=sle_chan_out
+                    )
+
+            layer = nn.ModuleList([
+                nn.Sequential(
+                    upsample(),
+                    Blur(),
+                    nn.Conv2d(chan_in, chan_out * 2, 3, padding=1),
+                    norm_class(chan_out * 2),
+                    nn.GLU(dim=1)
+                ),
+                sle,
+                attn
+            ])
+            self.layers.append(layer)
+
+        self.out_conv = nn.Conv2d(features[-1], init_channel, 3, padding=1)
+
+    def forward(self, x):
+        x = rearrange(x, 'b c -> b c () ()')
+        x = self.initial_conv(x)
+        x = F.normalize(x, dim=1)
+
+        residuals = dict()
+
+        for (res, (up, sle, attn)) in zip(self.res_layers, self.layers):
+            if exists(attn):
+                x = attn(x) + x
+
+            x = up(x)
+
+            if exists(sle):
+                out_res = self.sle_map[res]
+                residual = sle(x)
+                residuals[out_res] = residual
+
+            next_res = res + 1
+            if next_res in residuals:
+                x = x * residuals[next_res]
+
+        return self.out_conv(x)
+
+
+class SimpleDecoder(nn.Module):
+    def __init__(
+            self,
+            *,
+            chan_in,
+            chan_out=3,
+            num_upsamples=4,
+    ):
+        super().__init__()
+
+        self.layers = nn.ModuleList([])
+        final_chan = chan_out
+        chans = chan_in
+
+        for ind in range(num_upsamples):
+            last_layer = ind == (num_upsamples - 1)
+            chan_out = chans if not last_layer else final_chan * 2
+            layer = nn.Sequential(
+                upsample(),
+                nn.Conv2d(chans, chan_out, 3, padding=1),
+                nn.GLU(dim=1)
+            )
+            self.layers.append(layer)
+            chans //= 2
+
+    def forward(self, x):
+        for layer in self.layers:
+            x = layer(x)
+        return x
+
+
+class Discriminator(nn.Module):
+    def __init__(
+            self,
+            *,
+            image_size,
+            fmap_max=512,
+            fmap_inverse_coef=12,
+            transparent=False,
+            greyscale=False,
+            disc_output_size=5,
+            attn_res_layers=[]
+    ):
+        super().__init__()
+        resolution = log2(image_size)
+        assert is_power_of_two(image_size), 'image size must be a power of 2'
+        assert disc_output_size in {1, 5}, 'discriminator output dimensions can only be 5x5 or 1x1'
+
+        resolution = int(resolution)
+
+        if transparent:
+            init_channel = 4
+        elif greyscale:
+            init_channel = 1
+        else:
+            init_channel = 3
+
+        num_non_residual_layers = max(0, int(resolution) - 8)
+        num_residual_layers = 8 - 3
+
+        non_residual_resolutions = range(min(8, resolution), 2, -1)
+        features = list(map(lambda n: (n, 2 ** (fmap_inverse_coef - n)), non_residual_resolutions))
+        features = list(map(lambda n: (n[0], min(n[1], fmap_max)), features))
+
+        if num_non_residual_layers == 0:
+            res, _ = features[0]
+            features[0] = (res, init_channel)
+
+        chan_in_out = list(zip(features[:-1], features[1:]))
+
+        self.non_residual_layers = nn.ModuleList([])
+        for ind in range(num_non_residual_layers):
+            first_layer = ind == 0
+            last_layer = ind == (num_non_residual_layers - 1)
+            chan_out = features[0][-1] if last_layer else init_channel
+
+            self.non_residual_layers.append(nn.Sequential(
+                Blur(),
+                nn.Conv2d(init_channel, chan_out, 4, stride=2, padding=1),
+                nn.LeakyReLU(0.1)
+            ))
+
+        self.residual_layers = nn.ModuleList([])
+
+        for (res, ((_, chan_in), (_, chan_out))) in zip(non_residual_resolutions, chan_in_out):
+            image_width = 2 ** res
+
+            attn = None
+            if image_width in attn_res_layers:
+                attn = PreNorm(chan_in, LinearAttention(chan_in))
+
+            self.residual_layers.append(nn.ModuleList([
+                SumBranches([
+                    nn.Sequential(
+                        Blur(),
+                        nn.Conv2d(chan_in, chan_out, 4, stride=2, padding=1),
+                        nn.LeakyReLU(0.1),
+                        nn.Conv2d(chan_out, chan_out, 3, padding=1),
+                        nn.LeakyReLU(0.1)
+                    ),
+                    nn.Sequential(
+                        Blur(),
+                        nn.AvgPool2d(2),
+                        nn.Conv2d(chan_in, chan_out, 1),
+                        nn.LeakyReLU(0.1),
+                    )
+                ]),
+                attn
+            ]))
+
+        last_chan = features[-1][-1]
+        if disc_output_size == 5:
+            self.to_logits = nn.Sequential(
+                nn.Conv2d(last_chan, last_chan, 1),
+                nn.LeakyReLU(0.1),
+                nn.Conv2d(last_chan, 1, 4)
+            )
+        elif disc_output_size == 1:
+            self.to_logits = nn.Sequential(
+                Blur(),
+                nn.Conv2d(last_chan, last_chan, 3, stride=2, padding=1),
+                nn.LeakyReLU(0.1),
+                nn.Conv2d(last_chan, 1, 4)
+            )
+
+        self.to_shape_disc_out = nn.Sequential(
+            nn.Conv2d(init_channel, 64, 3, padding=1),
+            Residual(PreNorm(64, LinearAttention(64))),
+            SumBranches([
+                nn.Sequential(
+                    Blur(),
+                    nn.Conv2d(64, 32, 4, stride=2, padding=1),
+                    nn.LeakyReLU(0.1),
+                    nn.Conv2d(32, 32, 3, padding=1),
+                    nn.LeakyReLU(0.1)
+                ),
+                nn.Sequential(
+                    Blur(),
+                    nn.AvgPool2d(2),
+                    nn.Conv2d(64, 32, 1),
+                    nn.LeakyReLU(0.1),
+                )
+            ]),
+            Residual(PreNorm(32, LinearAttention(32))),
+            nn.AdaptiveAvgPool2d((4, 4)),
+            nn.Conv2d(32, 1, 4)
+        )
+
+        self.decoder1 = SimpleDecoder(chan_in=last_chan, chan_out=init_channel)
+        self.decoder2 = SimpleDecoder(chan_in=features[-2][-1], chan_out=init_channel) if resolution >= 9 else None
+
+    def forward(self, x, calc_aux_loss=False):
+        orig_img = x
+
+        for layer in self.non_residual_layers:
+            x = layer(x)
+
+        layer_outputs = []
+
+        for (net, attn) in self.residual_layers:
+            if exists(attn):
+                x = attn(x) + x
+
+            x = net(x)
+            layer_outputs.append(x)
+
+        out = self.to_logits(x).flatten(1)
+
+        img_32x32 = F.interpolate(orig_img, size=(32, 32))
+        out_32x32 = self.to_shape_disc_out(img_32x32)
+
+        if not calc_aux_loss:
+            return out, out_32x32, None
+
+        # self-supervised auto-encoding loss
+
+        layer_8x8 = layer_outputs[-1]
+        layer_16x16 = layer_outputs[-2]
+
+        recon_img_8x8 = self.decoder1(layer_8x8)
+
+        aux_loss = F.mse_loss(
+            recon_img_8x8,
+            F.interpolate(orig_img, size=recon_img_8x8.shape[2:])
+        )
+
+        if exists(self.decoder2):
+            select_random_quadrant = lambda rand_quadrant, img: \
+                rearrange(img, 'b c (m h) (n w) -> (m n) b c h w', m=2, n=2)[rand_quadrant]
+            crop_image_fn = partial(select_random_quadrant, floor(random() * 4))
+            img_part, layer_16x16_part = map(crop_image_fn, (orig_img, layer_16x16))
+
+            recon_img_16x16 = self.decoder2(layer_16x16_part)
+
+            aux_loss_16x16 = F.mse_loss(
+                recon_img_16x16,
+                F.interpolate(img_part, size=recon_img_16x16.shape[2:])
+            )
+
+            aux_loss = aux_loss + aux_loss_16x16
+
+        return out, out_32x32, aux_loss
+
+
+class LightweightGAN(nn.Module, HugGANModelHubMixin):
+    def __init__(
+            self,
+            *,
+            latent_dim,
+            image_size,
+            optimizer="adam",
+            fmap_max=512,
+            fmap_inverse_coef=12,
+            transparent=False,
+            greyscale=False,
+            disc_output_size=5,
+            attn_res_layers=[],
+            freq_chan_attn=False,
+            ttur_mult=1.,
+            lr=2e-4,
+    ):
+        super().__init__()
+
+        self.config = {
+            'latent_dim': latent_dim,
+            'image_size': image_size,
+            'optimizer': optimizer,
+            'fmap_max': fmap_max,
+            'fmap_inverse_coef': fmap_inverse_coef,
+            'transparent': transparent,
+            'greyscale': greyscale,
+            'disc_output_size': disc_output_size,
+            'attn_res_layers': attn_res_layers,
+            'freq_chan_attn': freq_chan_attn,
+            'ttur_mult': ttur_mult,
+            'lr': lr
+        }
+
+        self.latent_dim = latent_dim
+        self.image_size = image_size
+
+        G_kwargs = dict(
+            image_size=image_size,
+            latent_dim=latent_dim,
+            fmap_max=fmap_max,
+            fmap_inverse_coef=fmap_inverse_coef,
+            transparent=transparent,
+            greyscale=greyscale,
+            attn_res_layers=attn_res_layers,
+            freq_chan_attn=freq_chan_attn
+        )
+
+        self.G = Generator(**G_kwargs)
+
+        self.D = Discriminator(
+            image_size=image_size,
+            fmap_max=fmap_max,
+            fmap_inverse_coef=fmap_inverse_coef,
+            transparent=transparent,
+            greyscale=greyscale,
+            attn_res_layers=attn_res_layers,
+            disc_output_size=disc_output_size
+        )
+
+        self.ema_updater = EMA(0.995)
+        self.GE = Generator(**G_kwargs)
+        set_requires_grad(self.GE, False)
+
+        if optimizer == "adam":
+            self.G_opt = Adam(self.G.parameters(), lr=lr, betas=(0.5, 0.9))
+            self.D_opt = Adam(self.D.parameters(), lr=lr * ttur_mult, betas=(0.5, 0.9))
+        elif optimizer == "adabelief":
+            from adabelief_pytorch import AdaBelief
+
+            self.G_opt = AdaBelief(self.G.parameters(), lr=lr, betas=(0.5, 0.9))
+            self.D_opt = AdaBelief(self.D.parameters(), lr=lr * ttur_mult, betas=(0.5, 0.9))
+        else:
+            assert False, "No valid optimizer is given"
+
+        self.apply(self._init_weights)
+        self.reset_parameter_averaging()
+
+        self.D_aug = AugWrapper(self.D, image_size)
+
+    def _init_weights(self, m):
+        if type(m) in {nn.Conv2d, nn.Linear}:
+            nn.init.kaiming_normal_(m.weight, a=0, mode='fan_in', nonlinearity='leaky_relu')
+
+    def EMA(self):
+        def update_moving_average(ma_model, current_model):
+            for current_params, ma_params in zip(current_model.parameters(), ma_model.parameters()):
+                old_weight, up_weight = ma_params.data, current_params.data
+                ma_params.data = self.ema_updater.update_average(old_weight, up_weight)
+
+            for current_buffer, ma_buffer in zip(current_model.buffers(), ma_model.buffers()):
+                new_buffer_value = self.ema_updater.update_average(ma_buffer, current_buffer)
+                ma_buffer.copy_(new_buffer_value)
+
+        update_moving_average(self.GE, self.G)
+
+    def reset_parameter_averaging(self):
+        self.GE.load_state_dict(self.G.state_dict())
+
+    def forward(self, x):
+        raise NotImplemented
+
+    def _save_pretrained(self, save_directory):
+        """
+        Overwrite this method in case you don't want to save complete model,
+        rather some specific layers
+        """
+        path = os.path.join(save_directory, PYTORCH_WEIGHTS_NAME)
+        model_to_save = self.module if hasattr(self, "module") else self
+
+        # We update this to be a dict containing 'GAN', as that's what is expected
+        torch.save({'GAN': model_to_save.state_dict()}, path)
+
+    @classmethod
+    def _from_pretrained(
+            cls,
+            model_id,
+            revision,
+            cache_dir,
+            force_download,
+            proxies,
+            resume_download,
+            local_files_only,
+            token,
+            map_location="cpu",
+            strict=False,
+            **model_kwargs,
+    ):
+        """
+        Overwrite this method in case you wish to initialize your model in a
+        different way.
+        """
+        map_location = torch.device(map_location)
+
+        if os.path.isdir(model_id):
+            print("Loading weights from local directory")
+            model_file = os.path.join(model_id, PYTORCH_WEIGHTS_NAME)
+        else:
+            model_file = hf_hub_download(
+                repo_id=model_id,
+                filename=PYTORCH_WEIGHTS_NAME,
+                revision=revision,
+                cache_dir=cache_dir,
+                force_download=force_download,
+                proxies=proxies,
+                resume_download=resume_download,
+                token=token,
+                local_files_only=local_files_only,
+            )
+
+        # We update here to directly unpack config
+        model = cls(**model_kwargs['config'])
+
+        state_dict = torch.load(model_file, map_location=map_location)
+        model.load_state_dict(state_dict["GAN"], strict=strict)
+        model.eval()
+
+        return model
+
+
+# trainer
+
+class Trainer():
+    def __init__(
+            self,
+            dataset_name="huggan/CelebA-faces",
+            name='default',
+            results_dir='results',
+            models_dir='models',
+            base_dir='./',
+            optimizer='adam',
+            latent_dim=256,
+            image_size=128,
+            num_image_tiles=8,
+            fmap_max=512,
+            transparent=False,
+            greyscale=False,
+            batch_size=4,
+            gp_weight=10,
+            gradient_accumulate_every=1,
+            attn_res_layers=[],
+            freq_chan_attn=False,
+            disc_output_size=5,
+            dual_contrast_loss=False,
+            antialias=False,
+            lr=2e-4,
+            lr_mlp=1.,
+            ttur_mult=1.,
+            save_every=10000,
+            evaluate_every=1000,
+            aug_prob=None,
+            aug_types=['translation', 'cutout'],
+            dataset_aug_prob=0.,
+            calculate_fid_every=None,
+            calculate_fid_num_images=12800,
+            clear_fid_cache=False,
+            log=False,
+            cpu=False,
+            mixed_precision="no",
+            wandb=False,
+            push_to_hub=False,
+            organization_name=None,
+            *args,
+            **kwargs
+    ):
+        self.GAN_params = [args, kwargs]
+        self.GAN = None
+
+        self.dataset_name = dataset_name
+
+        self.name = name
+
+        base_dir = Path(base_dir)
+        self.base_dir = base_dir
+        self.results_dir = base_dir / results_dir
+        self.models_dir = base_dir / models_dir
+        self.fid_dir = base_dir / 'fid' / name
+
+        # Note - in original repo config is private - ".config.json", but here, we make it public
+        self.config_path = self.models_dir / name / 'config.json'
+
+        assert is_power_of_two(image_size), 'image size must be a power of 2 (64, 128, 256, 512, 1024)'
+        assert all(map(is_power_of_two,
+                       attn_res_layers)), 'resolution layers of attention must all be powers of 2 (16, 32, 64, 128, 256, 512)'
+
+        assert not (
+                dual_contrast_loss and disc_output_size > 1), 'discriminator output size cannot be greater than 1 if using dual contrastive loss'
+
+        self.image_size = image_size
+        self.num_image_tiles = num_image_tiles
+
+        self.latent_dim = latent_dim
+        self.fmap_max = fmap_max
+        self.transparent = transparent
+        self.greyscale = greyscale
+
+        assert (int(self.transparent) + int(self.greyscale)) < 2, 'you can only set either transparency or greyscale'
+
+        self.aug_prob = aug_prob
+        self.aug_types = aug_types
+
+        self.lr = lr
+        self.optimizer = optimizer
+        self.ttur_mult = ttur_mult
+        self.batch_size = batch_size
+        self.gradient_accumulate_every = gradient_accumulate_every
+
+        self.gp_weight = gp_weight
+
+        self.evaluate_every = evaluate_every
+        self.save_every = save_every
+        self.steps = 0
+
+        self.attn_res_layers = attn_res_layers
+        self.freq_chan_attn = freq_chan_attn
+
+        self.disc_output_size = disc_output_size
+        self.antialias = antialias
+
+        self.dual_contrast_loss = dual_contrast_loss
+
+        self.d_loss = 0
+        self.g_loss = 0
+        self.last_gp_loss = None
+        self.last_recon_loss = None
+        self.last_fid = None
+
+        self.init_folders()
+
+        self.loader = None
+        self.dataset_aug_prob = dataset_aug_prob
+
+        self.calculate_fid_every = calculate_fid_every
+        self.calculate_fid_num_images = calculate_fid_num_images
+        self.clear_fid_cache = clear_fid_cache
+
+        self.syncbatchnorm = torch.cuda.device_count() > 1 and not cpu
+
+        self.cpu = cpu
+        self.mixed_precision = mixed_precision
+
+        self.wandb = wandb
+
+        self.push_to_hub = push_to_hub
+        self.organization_name = organization_name
+        self.repo_name = get_full_repo_name(self.name, self.organization_name)
+        if self.push_to_hub:
+            self.repo_url = create_repo(self.repo_name, exist_ok=True)
+
+    @property
+    def image_extension(self):
+        return 'jpg' if not self.transparent else 'png'
+
+    @property
+    def checkpoint_num(self):
+        return floor(self.steps // self.save_every)
+
+    def init_GAN(self):
+        args, kwargs = self.GAN_params
+
+        # set some global variables before instantiating GAN
+
+        global norm_class
+        global Blur
+
+        norm_class = nn.SyncBatchNorm if self.syncbatchnorm else nn.BatchNorm2d
+        Blur = nn.Identity if not self.antialias else Fuzziness
+
+        # instantiate GAN
+
+        self.GAN = LightweightGAN(
+            optimizer=self.optimizer,
+            lr=self.lr,
+            latent_dim=self.latent_dim,
+            attn_res_layers=self.attn_res_layers,
+            freq_chan_attn=self.freq_chan_attn,
+            image_size=self.image_size,
+            ttur_mult=self.ttur_mult,
+            fmap_max=self.fmap_max,
+            disc_output_size=self.disc_output_size,
+            transparent=self.transparent,
+            greyscale=self.greyscale,
+            *args,
+            **kwargs
+        )
+
+    def write_config(self):
+        self.config_path.write_text(json.dumps(self.config()))
+
+    def load_config(self):
+        config = self.config() if not self.config_path.exists() else json.loads(self.config_path.read_text())
+        self.image_size = config['image_size']
+        self.transparent = config['transparent']
+        self.syncbatchnorm = config['syncbatchnorm']
+        self.disc_output_size = config['disc_output_size']
+        self.greyscale = config.pop('greyscale', False)
+        self.attn_res_layers = config.pop('attn_res_layers', [])
+        self.freq_chan_attn = config.pop('freq_chan_attn', False)
+        self.optimizer = config.pop('optimizer', 'adam')
+        self.fmap_max = config.pop('fmap_max', 512)
+        del self.GAN
+        self.init_GAN()
+
+    def config(self):
+        return {
+            'image_size': self.image_size,
+            'transparent': self.transparent,
+            'greyscale': self.greyscale,
+            'syncbatchnorm': self.syncbatchnorm,
+            'disc_output_size': self.disc_output_size,
+            'optimizer': self.optimizer,
+            'attn_res_layers': self.attn_res_layers,
+            'freq_chan_attn': self.freq_chan_attn
+        }
+
+    def set_data_src(self):
+        # start of using HuggingFace dataset
+        dataset = load_dataset(self.dataset_name)
+
+        if self.transparent:
+            num_channels = 4
+            pillow_mode = 'RGBA'
+            expand_fn = expand_greyscale(self.transparent)
+        elif self.greyscale:
+            num_channels = 1
+            pillow_mode = 'L'
+            expand_fn = identity()
+        else:
+            num_channels = 3
+            pillow_mode = 'RGB'
+            expand_fn = expand_greyscale(self.transparent)
+
+        convert_image_fn = partial(convert_image_to, pillow_mode)
+
+        transform = transforms.Compose([
+            transforms.Lambda(convert_image_fn),
+            transforms.Lambda(partial(resize_to_minimum_size, self.image_size)),
+            transforms.Resize(self.image_size),
+            RandomApply(0., transforms.RandomResizedCrop(self.image_size, scale=(0.5, 1.0), ratio=(0.98, 1.02)),
+                        transforms.CenterCrop(self.image_size)),
+            transforms.ToTensor(),
+            transforms.Lambda(expand_fn)
+        ])
+
+        def transform_images(examples):
+            transformed_images = [transform(image.convert("RGB")) for image in examples["image"]]
+
+            examples["image"] = torch.stack(transformed_images)
+
+            return examples
+
+        transformed_dataset = dataset.with_transform(transform_images)
+
+        per_device_batch_size = math.ceil(self.batch_size / self.accelerator.num_processes)
+        dataloader = DataLoader(transformed_dataset["train"], per_device_batch_size, sampler=None, shuffle=False,
+                                drop_last=True, pin_memory=True)
+        num_samples = len(transformed_dataset)
+        ## end of HuggingFace dataset
+
+        # Note - in original repo, this is wrapped with cycle, but we will do that after accelerator prepares
+        self.loader = dataloader
+
+        # auto set augmentation prob for user if dataset is detected to be low
+        # num_samples = len(self.dataset)
+        if not exists(self.aug_prob) and num_samples < 1e5:
+            self.aug_prob = min(0.5, (1e5 - num_samples) * 3e-6)
+            print(f'autosetting augmentation probability to {round(self.aug_prob * 100)}%')
+
+    def init_accelerator(self):
+        # Initialize the accelerator. We will let the accelerator handle device placement.
+        ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
+        self.accelerator = Accelerator(kwargs_handlers=[ddp_kwargs], mixed_precision=self.mixed_precision, cpu=self.cpu)
+        
+        if self.accelerator.is_local_main_process:
+            # set up Weights and Biases if requested
+            if self.wandb:
+                import wandb
+
+                wandb.init(project=str(self.results_dir).split("/")[-1])
+
+        if not exists(self.GAN):
+            self.init_GAN()
+
+        G = self.GAN.G
+        D = self.GAN.D
+        D_aug = self.GAN.D_aug
+
+        # discriminator loss fn
+
+        self.set_data_src()
+
+        # prepare
+        G, D, D_aug, self.GAN.D_opt, self.GAN.G_opt, self.loader = self.accelerator.prepare(G, D, D_aug, self.GAN.D_opt,
+                                                                                            self.GAN.G_opt, self.loader)
+        self.loader = cycle(self.loader)
+
+        return G, D, D_aug
+
+    def train(self, G, D, D_aug):
+        assert exists(self.loader), 'You must first initialize the data source with `.set_data_src(<folder of images>)`'
+
+        self.GAN.train()
+        total_disc_loss = torch.zeros([], device=self.accelerator.device)
+        total_gen_loss = torch.zeros([], device=self.accelerator.device)
+
+        batch_size = math.ceil(self.batch_size / self.accelerator.num_processes)
+
+        image_size = self.GAN.image_size
+        latent_dim = self.GAN.latent_dim
+
+        aug_prob = default(self.aug_prob, 0)
+        aug_types = self.aug_types
+        aug_kwargs = {'prob': aug_prob, 'types': aug_types}
+
+        apply_gradient_penalty = self.steps % 4 == 0
+
+        # discriminator loss fn
+
+        if self.dual_contrast_loss:
+            D_loss_fn = dual_contrastive_loss
+        else:
+            D_loss_fn = hinge_loss
+
+        # train discriminator
+
+        self.GAN.D_opt.zero_grad()
+        for i in range(self.gradient_accumulate_every):
+            latents = torch.randn(batch_size, latent_dim, device=self.accelerator.device)
+            image_batch = next(self.loader)["image"]
+            image_batch.requires_grad_()
+
+            with torch.no_grad():
+                generated_images = G(latents)
+
+            fake_output, fake_output_32x32, _ = D_aug(generated_images, detach=True, **aug_kwargs)
+
+            real_output, real_output_32x32, real_aux_loss = D_aug(image_batch, calc_aux_loss=True, **aug_kwargs)
+
+            real_output_loss = real_output
+            fake_output_loss = fake_output
+
+            divergence = D_loss_fn(real_output_loss, fake_output_loss)
+            divergence_32x32 = D_loss_fn(real_output_32x32, fake_output_32x32)
+            disc_loss = divergence + divergence_32x32
+
+            aux_loss = real_aux_loss
+            disc_loss = disc_loss + aux_loss
+
+            if apply_gradient_penalty:
+                outputs = [real_output, real_output_32x32]
+                if self.accelerator.scaler is not None:
+                    outputs = list(map(self.accelerator.scaler.scale, outputs))
+
+                scaled_gradients = torch_grad(outputs=outputs, inputs=image_batch,
+                                              grad_outputs=list(
+                                                  map(lambda t: torch.ones(t.size(), device=self.accelerator.device),
+                                                      outputs)),
+                                              create_graph=True, retain_graph=True, only_inputs=True)[0]
+
+                inv_scale = 1.
+                if self.accelerator.scaler is not None:
+                    inv_scale = safe_div(1., self.accelerator.scaler.get_scale())
+
+                if inv_scale != float('inf'):
+                    gradients = scaled_gradients * inv_scale
+
+                    gradients = gradients.reshape(batch_size, -1)
+                    gp = self.gp_weight * ((gradients.norm(2, dim=1) - 1) ** 2).mean()
+
+                    if not torch.isnan(gp):
+                        disc_loss = disc_loss + gp
+                        self.last_gp_loss = gp.clone().detach().item()
+
+            # divide loss by gradient accumulation steps since gradients
+            # are accumulated for multiple backward passes in PyTorch
+            disc_loss = disc_loss / self.gradient_accumulate_every
+
+            disc_loss.register_hook(raise_if_nan)
+            self.accelerator.backward(disc_loss)
+            total_disc_loss += divergence
+
+        self.last_recon_loss = aux_loss.item()
+        self.d_loss = float(total_disc_loss.item() / self.gradient_accumulate_every)
+        self.GAN.D_opt.step()
+
+        # generator loss fn
+
+        if self.dual_contrast_loss:
+            G_loss_fn = dual_contrastive_loss
+            G_requires_calc_real = True
+        else:
+            G_loss_fn = gen_hinge_loss
+            G_requires_calc_real = False
+
+        # train generator
+
+        self.GAN.G_opt.zero_grad()
+
+        for i in range(self.gradient_accumulate_every):
+            latents = torch.randn(batch_size, latent_dim, device=self.accelerator.device)
+
+            if G_requires_calc_real:
+                image_batch = next(self.loader)["image"]
+                image_batch.requires_grad_()
+
+            generated_images = G(latents)
+
+            fake_output, fake_output_32x32, _ = D_aug(generated_images, **aug_kwargs)
+            real_output, real_output_32x32, _ = D_aug(image_batch, **aug_kwargs) if G_requires_calc_real else (
+                None, None, None)
+
+            loss = G_loss_fn(fake_output, real_output)
+            loss_32x32 = G_loss_fn(fake_output_32x32, real_output_32x32)
+
+            gen_loss = loss + loss_32x32
+
+            gen_loss = gen_loss / self.gradient_accumulate_every
+
+            gen_loss.register_hook(raise_if_nan)
+            self.accelerator.backward(gen_loss)
+            total_gen_loss += loss
+
+            # divide loss by gradient accumulation steps since gradients
+        # are accumulated for multiple backward passes in PyTorch
+        self.g_loss = float(total_gen_loss.item() / self.gradient_accumulate_every)
+        self.GAN.G_opt.step()
+
+        # calculate moving averages
+        if self.accelerator.is_main_process and self.steps % 10 == 0 and self.steps > 20000:
+            self.GAN.EMA()
+
+        if self.accelerator.is_main_process and self.steps <= 25000 and self.steps % 1000 == 2:
+            self.GAN.reset_parameter_averaging()
+
+        # save from NaN errors
+
+        if any(torch.isnan(l) for l in (total_gen_loss, total_disc_loss)):
+            print(f'NaN detected for generator or discriminator. Loading from checkpoint #{self.checkpoint_num}')
+            self.load(self.checkpoint_num)
+            raise NanException
+
+        del total_disc_loss
+        del total_gen_loss
+
+        # periodically save results
+
+        if self.accelerator.is_main_process:
+            if self.steps % self.save_every == 0:
+                self.save(self.checkpoint_num)
+
+                if self.push_to_hub:
+                    with tempfile.TemporaryDirectory() as temp_dir:
+                        self.GAN.push_to_hub(temp_dir, self.repo_url, config=self.GAN.config, skip_lfs_files=True)
+
+            if self.steps % self.evaluate_every == 0 or (self.steps % 100 == 0 and self.steps < 20000):
+                self.evaluate(floor(self.steps / self.evaluate_every), num_image_tiles=self.num_image_tiles)
+
+            if exists(self.calculate_fid_every) and self.steps % self.calculate_fid_every == 0 and self.steps != 0:
+                num_batches = math.ceil(self.calculate_fid_num_images / self.batch_size)
+                fid = self.calculate_fid(num_batches)
+                self.last_fid = fid
+
+                with open(str(self.results_dir / self.name / f'fid_scores.txt'), 'a') as f:
+                    f.write(f'{self.steps},{fid}\n')
+
+        self.steps += 1
+
+    @torch.no_grad()
+    def evaluate(self, num=0, num_image_tiles=4):
+        self.GAN.eval()
+
+        ext = self.image_extension
+        num_rows = num_image_tiles
+
+        latent_dim = self.GAN.latent_dim
+        image_size = self.GAN.image_size
+
+        # latents and noise
+
+        latents = torch.randn(num_rows ** 2, latent_dim, device=self.accelerator.device)
+
+        # regular
+
+        generated_images = self.generate_(self.GAN.G, latents)
+        file_name = str(self.results_dir / self.name / f'{str(num)}.{ext}')
+        save_image(generated_images, file_name, nrow=num_rows)
+
+        # moving averages
+
+        generated_images = self.generate_(self.GAN.GE.to(self.accelerator.device), latents)
+        file_name_ema = str(self.results_dir / self.name / f'{str(num)}-ema.{ext}')
+        save_image(generated_images, file_name_ema, nrow=num_rows)
+
+        if self.accelerator.is_local_main_process and self.wandb:
+            import wandb
+
+            wandb.log({'generated_examples': wandb.Image(str(file_name))})
+            wandb.log({'generated_examples_ema': wandb.Image(str(file_name_ema))})
+
+    @torch.no_grad()
+    def generate(self, num=0, num_image_tiles=4, checkpoint=None, types=['default', 'ema']):
+        self.GAN.eval()
+
+        latent_dim = self.GAN.latent_dim
+        dir_name = self.name + str('-generated-') + str(checkpoint)
+        dir_full = Path().absolute() / self.results_dir / dir_name
+        ext = self.image_extension
+
+        if not dir_full.exists():
+            os.mkdir(dir_full)
+
+        # regular
+        if 'default' in types:
+            for i in tqdm(range(num_image_tiles), desc='Saving generated default images'):
+                latents = torch.randn(1, latent_dim, device=self.accelerator.device)
+                generated_image = self.generate_(self.GAN.G, latents)
+                path = str(self.results_dir / dir_name / f'{str(num)}-{str(i)}.{ext}')
+                save_image(generated_image[0], path, nrow=1)
+
+        # moving averages
+        if 'ema' in types:
+            for i in tqdm(range(num_image_tiles), desc='Saving generated EMA images'):
+                latents = torch.randn(1, latent_dim, device=self.accelerator.device)
+                generated_image = self.generate_(self.GAN.GE, latents)
+                path = str(self.results_dir / dir_name / f'{str(num)}-{str(i)}-ema.{ext}')
+                save_image(generated_image[0], path, nrow=1)
+
+        return dir_full
+
+    @torch.no_grad()
+    def show_progress(self, num_images=4, types=['default', 'ema']):
+        checkpoints = self.get_checkpoints()
+        assert exists(checkpoints), 'cannot find any checkpoints to create a training progress video for'
+
+        dir_name = self.name + str('-progress')
+        dir_full = Path().absolute() / self.results_dir / dir_name
+        ext = self.image_extension
+        latents = None
+
+        zfill_length = math.ceil(math.log10(len(checkpoints)))
+
+        if not dir_full.exists():
+            os.mkdir(dir_full)
+
+        for checkpoint in tqdm(checkpoints, desc='Generating progress images'):
+            self.load(checkpoint, print_version=False)
+            self.GAN.eval()
+
+            if checkpoint == 0:
+                latents = torch.randn(num_images, self.GAN.latent_dim, self.accelerator.device)
+
+            # regular
+            if 'default' in types:
+                generated_image = self.generate_(self.GAN.G, latents)
+                path = str(self.results_dir / dir_name / f'{str(checkpoint).zfill(zfill_length)}.{ext}')
+                save_image(generated_image, path, nrow=num_images)
+
+            # moving averages
+            if 'ema' in types:
+                generated_image = self.generate_(self.GAN.GE, latents)
+                path = str(self.results_dir / dir_name / f'{str(checkpoint).zfill(zfill_length)}-ema.{ext}')
+                save_image(generated_image, path, nrow=num_images)
+
+    @torch.no_grad()
+    def calculate_fid(self, num_batches):
+        from pytorch_fid import fid_score
+        real_path = self.fid_dir / 'real'
+        fake_path = self.fid_dir / 'fake'
+
+        # remove any existing files used for fid calculation and recreate directories
+        if not real_path.exists() or self.clear_fid_cache:
+            rmtree(real_path, ignore_errors=True)
+            os.makedirs(real_path)
+
+            for batch_num in tqdm(range(num_batches), desc='calculating FID - saving reals'):
+                real_batch = next(self.loader)["image"]
+                for k, image in enumerate(real_batch.unbind(0)):
+                    ind = k + batch_num * self.batch_size
+                    save_image(image, real_path / f'{ind}.png')
+
+        # generate a bunch of fake images in results / name / fid_fake
+        rmtree(fake_path, ignore_errors=True)
+        os.makedirs(fake_path)
+
+        self.GAN.eval()
+        ext = self.image_extension
+
+        latent_dim = self.GAN.latent_dim
+        image_size = self.GAN.image_size
+
+        for batch_num in tqdm(range(num_batches), desc='calculating FID - saving generated'):
+            # latents and noise
+            latents = torch.randn(self.batch_size, latent_dim, device=self.accelerator.device)
+
+            # moving averages
+            generated_images = self.generate_(self.GAN.GE, latents)
+
+            for j, image in enumerate(generated_images.unbind(0)):
+                ind = j + batch_num * self.batch_size
+                save_image(image, str(fake_path / f'{str(ind)}-ema.{ext}'))
+
+        return fid_score.calculate_fid_given_paths([str(real_path), str(fake_path)], 256, latents.device, 2048)
+
+    @torch.no_grad()
+    def generate_(self, G, style, num_image_tiles=8):
+        generated_images = evaluate_in_chunks(self.batch_size, G, style)
+        return generated_images.clamp_(0., 1.)
+
+    @torch.no_grad()
+    def generate_interpolation(self, num=0, num_image_tiles=8, num_steps=100, save_frames=False):
+        self.GAN.eval()
+        ext = self.image_extension
+        num_rows = num_image_tiles
+
+        latent_dim = self.GAN.latent_dim
+        image_size = self.GAN.image_size
+
+        # latents and noise
+        latents_low = torch.randn(num_rows ** 2, latent_dim, device=self.accelerator.device)
+        latents_high = torch.randn(num_rows ** 2, latent_dim, device=self.accelerator.device)
+
+        ratios = torch.linspace(0., 8., num_steps)
+
+        frames = []
+        for ratio in tqdm(ratios):
+            interp_latents = slerp(ratio, latents_low, latents_high)
+            generated_images = self.generate_(self.GAN.GE, interp_latents)
+            images_grid = torchvision.utils.make_grid(generated_images, nrow=num_rows)
+            pil_image = transforms.ToPILImage()(images_grid.cpu())
+
+            if self.transparent:
+                background = Image.new('RGBA', pil_image.size, (255, 255, 255))
+                pil_image = Image.alpha_composite(background, pil_image)
+
+            frames.append(pil_image)
+
+        frames[0].save(str(self.results_dir / self.name / f'{str(num)}.gif'), save_all=True, append_images=frames[1:],
+                       duration=80, loop=0, optimize=True)
+
+        if save_frames:
+            folder_path = (self.results_dir / self.name / f'{str(num)}')
+            folder_path.mkdir(parents=True, exist_ok=True)
+            for ind, frame in enumerate(frames):
+                frame.save(str(folder_path / f'{str(ind)}.{ext}'))
+
+    def print_log(self):
+        data = [
+            ('G', self.g_loss),
+            ('D', self.d_loss),
+            ('GP', self.last_gp_loss),
+            ('SS', self.last_recon_loss),
+            ('FID', self.last_fid)
+        ]
+
+        data = [d for d in data if exists(d[1])]
+        log = ' | '.join(map(lambda n: f'{n[0]}: {n[1]:.2f}', data))
+        print(log)
+
+        if self.accelerator.is_local_main_process:
+            log_dict = {v[0]: v[1] for v in data}
+            if self.wandb:
+                import wandb
+
+                wandb.log(log_dict)
+
+    def model_name(self, num):
+        return str(self.models_dir / self.name / f'model_{num}.pt')
+
+    def init_folders(self):
+        (self.results_dir / self.name).mkdir(parents=True, exist_ok=True)
+        (self.models_dir / self.name).mkdir(parents=True, exist_ok=True)
+
+    def clear(self):
+        rmtree(str(self.models_dir / self.name), True)
+        rmtree(str(self.results_dir / self.name), True)
+        rmtree(str(self.fid_dir), True)
+        rmtree(str(self.config_path), True)
+        self.init_folders()
+
+    def save(self, num):
+        save_data = {
+            'GAN': self.GAN.state_dict(),
+        }
+
+        torch.save(save_data, self.model_name(num))
+        self.write_config()
+
+    def load(self, num=-1):
+        self.load_config()
+
+        name = num
+        if num == -1:
+            checkpoints = self.get_checkpoints()
+
+            if not exists(checkpoints):
+                return
+
+            name = checkpoints[-1]
+            print(f'continuing from previous epoch - {name}')
+
+        self.steps = name * self.save_every
+
+        load_data = torch.load(self.model_name(name))
+
+        try:
+            self.GAN.load_state_dict(load_data['GAN'])
+        except Exception as e:
+            print(
+                'unable to load save model. please try downgrading the package to the version specified by the saved model')
+            raise e
+
+    def get_checkpoints(self):
+        file_paths = [p for p in Path(self.models_dir / self.name).glob('model_*.pt')]
+        saved_nums = sorted(map(lambda x: int(x.stem.split('_')[1]), file_paths))
+
+        if len(saved_nums) == 0:
+            return None
+
+        return saved_nums

huggan/pytorch/metrics/README.md

@@ -0,0 +1,39 @@
+# GAN metrics
+
+In order to track progress 📈 in (un)conditional image generation, a few quantitative metrics have been proposed. Below, we explain the most popular ones. For a more extensive overview, we refer the reader to [Borji, 2021](https://arxiv.org/abs/2103.09396) - which is an up-to-date version of [Borji, 2018](https://arxiv.org/abs/1802.03446). The TLDR is that, despite the use of many popular metrics, objective and comprehensive evaluation of generative models is still an open problem 🤷‍♂️.
+
+Quantitative metrics are of course just a proxy of image quality. The most widely used (Inception Score and FID) have several drawbacks [Barratt et al., 2018](https://arxiv.org/abs/1801.01973), [Sajjadi et al., 2018](https://arxiv.org/abs/1806.00035), [Kynkäänniemi et al., 2019](https://arxiv.org/abs/1904.06991).
+
+## Inception score
+
+The Inception score was proposed in [Salimans et al., 2016](https://arxiv.org/abs/1606.03498). The authors used a pre-trained Inceptionv3 neural net to classify the images generated by a GAN, and computed a score based on the class probablities of the neural net. The authors claimed that the score correlates well with subjective human evaluation. For an extensive explanation of the metric (as well as an implementation in Numpy and Keras), we refer the reader to [this blog post](https://machinelearningmastery.com/how-to-implement-the-inception-score-from-scratch-for-evaluating-generated-images/#:~:text=The%20Inception%20Score%2C%20or%20IS%20for%20short%2C%20is%20an%20objective,Improved%20Techniques%20for%20Training%20GANs.%E2%80%9D).
+
+## Fréchet Inception Distance (FID)
+
+The FID metric was proposed in [Heusel et al., 2018](https://arxiv.org/abs/1706.08500), and is currently the most widely used metric for evaluating image generation. Rather than only evaluating the generated images (as the Inception score), the FID metric compares the generated images to real images.
+
+The Fréchet distance meaures the distance between 2 multivariate Gaussian distributions. What does that mean? Concretely, the FID metric uses a pre-trained neural network (the same one as the one of the Inception score, Inceptionv3), and first forwards both real and generated images through it in order to get feature maps. Next, one computes statistics (namely, the mean and standard deviation) of the feature maps for both distributions (generated and real images). Finally, the distance between both distributions is computed based on these statistics.
+
+The FID metric assumes that feature maps of a pre-trained neural net extracted on real vs. fake images should be similar (the authors argue that this is a good quantitative metric for assessing image quality, correlating well with human judgement). 
+
+An important disadvantage of the FID metric is that is has an issue of generalization; a model that simply memorizes the training data can obtain a perfect score on these metrics [Razavi et al., 2019](https://arxiv.org/abs/1906.00446).
+
+Variants have been proposed for other modalities, such as the Fréchet Audio Distance [Kilgour et al., 2018](https://arxiv.org/abs/1812.08466) and the Fréchet Video Distance [Unterthiner et al., 2018](https://arxiv.org/abs/1812.01717). 
+
+The official implementation is in Tensorflow and can be found [here](https://github.com/bioinf-jku/TTUR). A PyTorch implementation can be found [here](https://github.com/mseitzer/pytorch-fid).
+
+## Clean FID
+
+In 2021, a paper by [Parmar et al.](https://arxiv.org/abs/2104.11222) indicated that the FID metric is often poorly computed, due to incorrect implementations of low-level image preprocessing (such as resizing of images) in popular frameworks such as PyTorch and TensorFlow. This can produce widely different values for the FID metric.
+
+The official implementation of the cleaner FID version can be found [here](https://github.com/GaParmar/clean-fid).
+
+Note that FID has many, many other variants including spatial FID (sFID), class-aware FID (CAFD) and conditional FID, Fast FID, Memorization-informed FID (MiFID), Unbiased FID, etc.
+
+## Precision and Recall
+
+Despite the FID metric being popular and correlating well with human evaluation, [Sajjadi et al., 2018](https://arxiv.org/abs/1806.00035) pointed out that, due to the fact that the FID score is just a scalar number, it is unable to distinguish between different failure cases. Two generative models could obtain the same FID score while generating images that look entirely different. Hence, the authors proposed a novel approach, defining precision (P) and recall (R) for distributions.
+
+Precision measures the similarity of generated instances to the real ones and recall measures the ability of a generator to synthesize all instances found in the training set. Hence, precision measures the quality and recall the coverage.
+
+These metrics were then further improved by [Kynkäänniemi et al., 2019](https://arxiv.org/abs/1904.06991).

huggan/pytorch/metrics/fid_score.py

@@ -0,0 +1,80 @@
+# sources: 
+# https://www.kaggle.com/code/ibtesama/gan-in-pytorch-with-fid/notebook
+# https://github.com/mseitzer/pytorch-fid/blob/master/src/pytorch_fid/fid_score.py
+
+import numpy as np
+from scipy import linalg
+from torch.nn.functional import adaptive_avg_pool2d
+
+
+def calculate_activation_statistics(images, model, batch_size=128, dims=2048):
+    model.eval()
+    act = np.empty((len(images), dims))
+
+    batch = images
+    pred = model(batch)[0]
+
+    # If model output is not scalar, apply global spatial average pooling.
+    # This happens if you choose a dimensionality not equal 2048.
+    if pred.size(2) != 1 or pred.size(3) != 1:
+        pred = adaptive_avg_pool2d(pred, output_size=(1, 1))
+
+    act = pred.cpu().data.numpy().reshape(pred.size(0), -1)
+    
+    mu = np.mean(act, axis=0)
+    sigma = np.cov(act, rowvar=False)
+    return mu, sigma
+
+
+def calculate_frechet_distance(mu1, sigma1, mu2, sigma2, eps=1e-6):
+    """Numpy implementation of the Frechet Distance.
+    The Frechet distance between two multivariate Gaussians X_1 ~ N(mu_1, C_1)
+    and X_2 ~ N(mu_2, C_2) is
+            d^2 = ||mu_1 - mu_2||^2 + Tr(C_1 + C_2 - 2*sqrt(C_1*C_2)).
+    """
+
+    mu1 = np.atleast_1d(mu1)
+    mu2 = np.atleast_1d(mu2)
+
+    sigma1 = np.atleast_2d(sigma1)
+    sigma2 = np.atleast_2d(sigma2)
+
+    assert mu1.shape == mu2.shape, \
+        'Training and test mean vectors have different lengths'
+    assert sigma1.shape == sigma2.shape, \
+        'Training and test covariances have different dimensions'
+
+    diff = mu1 - mu2
+
+    
+    covmean, _ = linalg.sqrtm(sigma1.dot(sigma2), disp=False)
+    if not np.isfinite(covmean).all():
+        msg = ('fid calculation produces singular product; '
+               'adding %s to diagonal of cov estimates') % eps
+        print(msg)
+        offset = np.eye(sigma1.shape[0]) * eps
+        covmean = linalg.sqrtm((sigma1 + offset).dot(sigma2 + offset))
+
+    
+    if np.iscomplexobj(covmean):
+        if not np.allclose(np.diagonal(covmean).imag, 0, atol=1e-3):
+            m = np.max(np.abs(covmean.imag))
+            raise ValueError('Imaginary component {}'.format(m))
+        covmean = covmean.real
+
+    tr_covmean = np.trace(covmean)
+
+    return (diff.dot(diff) + np.trace(sigma1) +
+            np.trace(sigma2) - 2 * tr_covmean)
+
+
+def calculate_fretchet(images_real, images_fake, model):
+    """Calculate the fretched distance."""
+     
+    # calculate statistics (mean + std)
+    mu_1, std_1 = calculate_activation_statistics(images_real, model)
+    mu_2, std_2 = calculate_activation_statistics(images_fake, model)
+    
+    # compute distance
+    fid_value = calculate_frechet_distance(mu_1, std_1, mu_2, std_2)
+    return fid_value

huggan/pytorch/metrics/inception.py

@@ -0,0 +1,328 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision
+
+try:
+    from torchvision.models.utils import load_state_dict_from_url
+except ImportError:
+    from torch.utils.model_zoo import load_url as load_state_dict_from_url
+
+# Inception weights ported to Pytorch from
+# http://download.tensorflow.org/models/image/imagenet/inception-2015-12-05.tgz
+FID_WEIGHTS_URL = 'https://github.com/mseitzer/pytorch-fid/releases/download/fid_weights/pt_inception-2015-12-05-6726825d.pth'  # noqa: E501
+
+
+class InceptionV3(nn.Module):
+    """Pretrained InceptionV3 network returning feature maps"""
+
+    # Index of default block of inception to return,
+    # corresponds to output of final average pooling
+    DEFAULT_BLOCK_INDEX = 3
+
+    # Maps feature dimensionality to their output blocks indices
+    BLOCK_INDEX_BY_DIM = {
+        64: 0,   # First max pooling features
+        192: 1,  # Second max pooling featurs
+        768: 2,  # Pre-aux classifier features
+        2048: 3  # Final average pooling features
+    }
+
+    def __init__(self,
+                 output_blocks=(DEFAULT_BLOCK_INDEX,),
+                 resize_input=True,
+                 normalize_input=True,
+                 requires_grad=False,
+                 use_fid_inception=True):
+        """Build pretrained InceptionV3
+
+        Parameters
+        ----------
+        output_blocks : list of int
+            Indices of blocks to return features of. Possible values are:
+                - 0: corresponds to output of first max pooling
+                - 1: corresponds to output of second max pooling
+                - 2: corresponds to output which is fed to aux classifier
+                - 3: corresponds to output of final average pooling
+        resize_input : bool
+            If true, bilinearly resizes input to width and height 299 before
+            feeding input to model. As the network without fully connected
+            layers is fully convolutional, it should be able to handle inputs
+            of arbitrary size, so resizing might not be strictly needed
+        normalize_input : bool
+            If true, scales the input from range (0, 1) to the range the
+            pretrained Inception network expects, namely (-1, 1)
+        requires_grad : bool
+            If true, parameters of the model require gradients. Possibly useful
+            for finetuning the network
+        use_fid_inception : bool
+            If true, uses the pretrained Inception model used in Tensorflow's
+            FID implementation. If false, uses the pretrained Inception model
+            available in torchvision. The FID Inception model has different
+            weights and a slightly different structure from torchvision's
+            Inception model. If you want to compute FID scores, you are
+            strongly advised to set this parameter to true to get comparable
+            results.
+        """
+        super(InceptionV3, self).__init__()
+
+        self.resize_input = resize_input
+        self.normalize_input = normalize_input
+        self.output_blocks = sorted(output_blocks)
+        self.last_needed_block = max(output_blocks)
+
+        assert self.last_needed_block <= 3, \
+            'Last possible output block index is 3'
+
+        self.blocks = nn.ModuleList()
+
+        if use_fid_inception:
+            inception = fid_inception_v3()
+        else:
+            inception = _inception_v3(pretrained=True)
+
+        # Block 0: input to maxpool1
+        block0 = [
+            inception.Conv2d_1a_3x3,
+            inception.Conv2d_2a_3x3,
+            inception.Conv2d_2b_3x3,
+            nn.MaxPool2d(kernel_size=3, stride=2)
+        ]
+        self.blocks.append(nn.Sequential(*block0))
+
+        # Block 1: maxpool1 to maxpool2
+        if self.last_needed_block >= 1:
+            block1 = [
+                inception.Conv2d_3b_1x1,
+                inception.Conv2d_4a_3x3,
+                nn.MaxPool2d(kernel_size=3, stride=2)
+            ]
+            self.blocks.append(nn.Sequential(*block1))
+
+        # Block 2: maxpool2 to aux classifier
+        if self.last_needed_block >= 2:
+            block2 = [
+                inception.Mixed_5b,
+                inception.Mixed_5c,
+                inception.Mixed_5d,
+                inception.Mixed_6a,
+                inception.Mixed_6b,
+                inception.Mixed_6c,
+                inception.Mixed_6d,
+                inception.Mixed_6e,
+            ]
+            self.blocks.append(nn.Sequential(*block2))
+
+        # Block 3: aux classifier to final avgpool
+        if self.last_needed_block >= 3:
+            block3 = [
+                inception.Mixed_7a,
+                inception.Mixed_7b,
+                inception.Mixed_7c,
+                nn.AdaptiveAvgPool2d(output_size=(1, 1))
+            ]
+            self.blocks.append(nn.Sequential(*block3))
+
+        for param in self.parameters():
+            param.requires_grad = requires_grad
+
+    def forward(self, inp):
+        """Get Inception feature maps
+
+        Parameters
+        ----------
+        inp : torch.autograd.Variable
+            Input tensor of shape Bx3xHxW. Values are expected to be in
+            range (0, 1)
+
+        Returns
+        -------
+        List of torch.autograd.Variable, corresponding to the selected output
+        block, sorted ascending by index
+        """
+        outp = []
+        x = inp
+
+        if self.resize_input:
+            x = F.interpolate(x,
+                              size=(299, 299),
+                              mode='bilinear',
+                              align_corners=False)
+
+        if self.normalize_input:
+            x = 2 * x - 1  # Scale from range (0, 1) to range (-1, 1)
+
+        for idx, block in enumerate(self.blocks):
+            x = block(x)
+            if idx in self.output_blocks:
+                outp.append(x)
+
+            if idx == self.last_needed_block:
+                break
+
+        return outp
+
+
+def _inception_v3(*args, **kwargs):
+    """Wraps `torchvision.models.inception_v3`
+
+    Skips default weight inititialization if supported by torchvision version.
+    See https://github.com/mseitzer/pytorch-fid/issues/28.
+    """
+    try:
+        version = tuple(map(int, torchvision.__version__.split('.')[:2]))
+    except ValueError:
+        # Just a caution against weird version strings
+        version = (0,)
+
+    if version >= (0, 6):
+        kwargs['init_weights'] = False
+
+    return torchvision.models.inception_v3(*args, **kwargs)
+
+
+def fid_inception_v3():
+    """Build pretrained Inception model for FID computation
+
+    The Inception model for FID computation uses a different set of weights
+    and has a slightly different structure than torchvision's Inception.
+
+    This method first constructs torchvision's Inception and then patches the
+    necessary parts that are different in the FID Inception model.
+    """
+    inception = _inception_v3(num_classes=1008,
+                              aux_logits=False,
+                              pretrained=False)
+    inception.Mixed_5b = FIDInceptionA(192, pool_features=32)
+    inception.Mixed_5c = FIDInceptionA(256, pool_features=64)
+    inception.Mixed_5d = FIDInceptionA(288, pool_features=64)
+    inception.Mixed_6b = FIDInceptionC(768, channels_7x7=128)
+    inception.Mixed_6c = FIDInceptionC(768, channels_7x7=160)
+    inception.Mixed_6d = FIDInceptionC(768, channels_7x7=160)
+    inception.Mixed_6e = FIDInceptionC(768, channels_7x7=192)
+    inception.Mixed_7b = FIDInceptionE_1(1280)
+    inception.Mixed_7c = FIDInceptionE_2(2048)
+
+    state_dict = load_state_dict_from_url(FID_WEIGHTS_URL, progress=True)
+    inception.load_state_dict(state_dict)
+    return inception
+
+
+class FIDInceptionA(torchvision.models.inception.InceptionA):
+    """InceptionA block patched for FID computation"""
+    def __init__(self, in_channels, pool_features):
+        super(FIDInceptionA, self).__init__(in_channels, pool_features)
+
+    def forward(self, x):
+        branch1x1 = self.branch1x1(x)
+
+        branch5x5 = self.branch5x5_1(x)
+        branch5x5 = self.branch5x5_2(branch5x5)
+
+        branch3x3dbl = self.branch3x3dbl_1(x)
+        branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)
+        branch3x3dbl = self.branch3x3dbl_3(branch3x3dbl)
+
+        # Patch: Tensorflow's average pool does not use the padded zero's in
+        # its average calculation
+        branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1,
+                                   count_include_pad=False)
+        branch_pool = self.branch_pool(branch_pool)
+
+        outputs = [branch1x1, branch5x5, branch3x3dbl, branch_pool]
+        return torch.cat(outputs, 1)
+
+
+class FIDInceptionC(torchvision.models.inception.InceptionC):
+    """InceptionC block patched for FID computation"""
+    def __init__(self, in_channels, channels_7x7):
+        super(FIDInceptionC, self).__init__(in_channels, channels_7x7)
+
+    def forward(self, x):
+        branch1x1 = self.branch1x1(x)
+
+        branch7x7 = self.branch7x7_1(x)
+        branch7x7 = self.branch7x7_2(branch7x7)
+        branch7x7 = self.branch7x7_3(branch7x7)
+
+        branch7x7dbl = self.branch7x7dbl_1(x)
+        branch7x7dbl = self.branch7x7dbl_2(branch7x7dbl)
+        branch7x7dbl = self.branch7x7dbl_3(branch7x7dbl)
+        branch7x7dbl = self.branch7x7dbl_4(branch7x7dbl)
+        branch7x7dbl = self.branch7x7dbl_5(branch7x7dbl)
+
+        # Patch: Tensorflow's average pool does not use the padded zero's in
+        # its average calculation
+        branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1,
+                                   count_include_pad=False)
+        branch_pool = self.branch_pool(branch_pool)
+
+        outputs = [branch1x1, branch7x7, branch7x7dbl, branch_pool]
+        return torch.cat(outputs, 1)
+
+
+class FIDInceptionE_1(torchvision.models.inception.InceptionE):
+    """First InceptionE block patched for FID computation"""
+    def __init__(self, in_channels):
+        super(FIDInceptionE_1, self).__init__(in_channels)
+
+    def forward(self, x):
+        branch1x1 = self.branch1x1(x)
+
+        branch3x3 = self.branch3x3_1(x)
+        branch3x3 = [
+            self.branch3x3_2a(branch3x3),
+            self.branch3x3_2b(branch3x3),
+        ]
+        branch3x3 = torch.cat(branch3x3, 1)
+
+        branch3x3dbl = self.branch3x3dbl_1(x)
+        branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)
+        branch3x3dbl = [
+            self.branch3x3dbl_3a(branch3x3dbl),
+            self.branch3x3dbl_3b(branch3x3dbl),
+        ]
+        branch3x3dbl = torch.cat(branch3x3dbl, 1)
+
+        # Patch: Tensorflow's average pool does not use the padded zero's in
+        # its average calculation
+        branch_pool = F.avg_pool2d(x, kernel_size=3, stride=1, padding=1,
+                                   count_include_pad=False)
+        branch_pool = self.branch_pool(branch_pool)
+
+        outputs = [branch1x1, branch3x3, branch3x3dbl, branch_pool]
+        return torch.cat(outputs, 1)
+
+
+class FIDInceptionE_2(torchvision.models.inception.InceptionE):
+    """Second InceptionE block patched for FID computation"""
+    def __init__(self, in_channels):
+        super(FIDInceptionE_2, self).__init__(in_channels)
+
+    def forward(self, x):
+        branch1x1 = self.branch1x1(x)
+
+        branch3x3 = self.branch3x3_1(x)
+        branch3x3 = [
+            self.branch3x3_2a(branch3x3),
+            self.branch3x3_2b(branch3x3),
+        ]
+        branch3x3 = torch.cat(branch3x3, 1)
+
+        branch3x3dbl = self.branch3x3dbl_1(x)
+        branch3x3dbl = self.branch3x3dbl_2(branch3x3dbl)
+        branch3x3dbl = [
+            self.branch3x3dbl_3a(branch3x3dbl),
+            self.branch3x3dbl_3b(branch3x3dbl),
+        ]
+        branch3x3dbl = torch.cat(branch3x3dbl, 1)
+
+        # Patch: The FID Inception model uses max pooling instead of average
+        # pooling. This is likely an error in this specific Inception
+        # implementation, as other Inception models use average pooling here
+        # (which matches the description in the paper).
+        branch_pool = F.max_pool2d(x, kernel_size=3, stride=1, padding=1)
+        branch_pool = self.branch_pool(branch_pool)
+
+        outputs = [branch1x1, branch3x3, branch3x3dbl, branch_pool]
+        return torch.cat(outputs, 1)