enable model loading from hf hub

.gitignore

@@ -20,6 +20,7 @@ __pycache__/
 *.zip
 events.*
 
+/checkpoints/
 *.pkl
 *.h5
 *.dat

.ipynb_checkpoints/model_to_hf_hub-checkpoint.ipynb

@@ -0,0 +1,255 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "import huggingface_hub\n",
+    "import utils"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "525a0eaa021f4fdebd9138f4e7c5ab65",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "VBox(children=(HTML(value='<center> <img\\nsrc=https://huggingface.co/front/assets/huggingface_logo-noborder.sv…"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "huggingface_hub.notebook_login()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Building generator for model `stylegan_animeface512` ...\n",
+      "Finish building generator.\n",
+      "Loading checkpoint from `checkpoints/stylegan_animeface512.pth` ...\n",
+      "Finish loading checkpoint.\n"
+     ]
+    }
+   ],
+   "source": [
+    "animeface_model = utils.load_generator('stylegan_animeface512')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Cloning https://huggingface.co/johko/stylegan_animeface512 into local empty directory.\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "6e51c5ae4a504617aa0f1c1ac798ed15",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Upload file pytorch_model.bin:   0%|          | 32.0k/103M [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "To https://huggingface.co/johko/stylegan_animeface512\n",
+      "   750cd03..2841156  main -> main\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "'https://huggingface.co/johko/stylegan_animeface512/commit/2841156bad3c5a5f47f3edbf4a41880ea8fd3ad3'"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "animeface_model.push_to_hub(\"johko/stylegan_animeface512\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Building generator for model `pggan_celebahq1024` ...\n",
+      "Finish building generator.\n",
+      "Loading checkpoint from `checkpoints/pggan_celebahq1024.pth` ...\n",
+      "Finish loading checkpoint.\n"
+     ]
+    }
+   ],
+   "source": [
+    "celebhq_model = utils.load_generator(\"pggan_celebahq1024\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Cloning https://huggingface.co/johko/pggan-celebahq-1024 into local empty directory.\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "ef4086b23a654b079bd6a3678140c50d",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Upload file pytorch_model.bin:   0%|          | 32.0k/88.1M [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "To https://huggingface.co/johko/pggan-celebahq-1024\n",
+      "   780695e..278449f  main -> main\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "'https://huggingface.co/johko/pggan-celebahq-1024/commit/278449f8416d38a0233c980774528d32c4eee99c'"
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "celebhq_model.push_to_hub(\"johko/pggan-celebahq-1024\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Building generator for model `stylegan_car512` ...\n",
+      "Finish building generator.\n",
+      "Loading checkpoint from `checkpoints/stylegan_car512.pth` ...\n",
+      "Finish loading checkpoint.\n"
+     ]
+    }
+   ],
+   "source": [
+    "cars_model = utils.load_generator(\"stylegan_car512\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "cars_model.push_to_hub(\"johko/stylegan_car512\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "interpreter": {
+   "hash": "a8d699d01f596cc27ac2722fbc0550b939d217978c7e1ca888dca7ba146ee4bf"
+  },
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

app.py

@@ -16,7 +16,7 @@ from utils import factorize_weight
 @st.cache(allow_output_mutation=True, show_spinner=False)
 def get_model(model_name):
     """Gets model by name."""
-    return load_generator(model_name)
+    return load_generator(model_name, from_hf_hub=True)
 
 
 @st.cache(allow_output_mutation=True, show_spinner=False)
@@ -72,7 +72,7 @@ layer_idx = st.sidebar.selectbox(
 layers, boundaries, eigen_values = factorize_model(model, layer_idx)
 
 num_semantics = st.sidebar.number_input(
-    'Number of semantics', value=10, min_value=0, max_value=None, step=1)
+    'Number of semantics', value=5, min_value=0, max_value=None, step=1)
 steps = {sem_idx: 0 for sem_idx in range(num_semantics)}
 if gan_type == 'pggan':
     max_step = 5.0

interface.py

@@ -16,7 +16,7 @@ from utils import factorize_weight
 @st.cache(allow_output_mutation=True, show_spinner=False)
 def get_model(model_name):
     """Gets model by name."""
-    return load_generator(model_name)
+    return load_generator(model_name, from_hf_hub=True)
 
 
 @st.cache(allow_output_mutation=True, show_spinner=False)
@@ -27,7 +27,7 @@ def factorize_model(model, layer_idx):
 
 def sample(model, gan_type, num=1):
     """Samples latent codes."""
-    codes = torch.randn(num, model.z_space_dim).cuda()
+    codes = torch.randn(num, model.z_space_dim)
     if gan_type == 'pggan':
         codes = model.layer0.pixel_norm(codes)
     elif gan_type == 'stylegan':
@@ -63,8 +63,7 @@ def main():
 
     model_name = st.sidebar.selectbox(
         'Model to Interpret',
-        ['pggan_celebahq1024', 'stylegan_animeface512', 'stylegan_car512', 'stylegan_cat256'
-         ])
+        ['pggan_celebahq1024', 'stylegan_animeface512', 'stylegan_car512', 'stylegan_cat256',])
 
     model = get_model(model_name)
     gan_type = parse_gan_type(model)
@@ -74,7 +73,7 @@ def main():
     layers, boundaries, eigen_values = factorize_model(model, layer_idx)
 
     num_semantics = st.sidebar.number_input(
-        'Number of semantics', value=10, min_value=0, max_value=None, step=1)
+        'Number of semantics', value=5, min_value=0, max_value=None, step=1)
     steps = {sem_idx: 0 for sem_idx in range(num_semantics)}
     if gan_type == 'pggan':
         max_step = 5.0

model_to_hf_hub.ipynb

@@ -0,0 +1,297 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "import huggingface_hub\n",
+    "import utils"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "525a0eaa021f4fdebd9138f4e7c5ab65",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "VBox(children=(HTML(value='<center> <img\\nsrc=https://huggingface.co/front/assets/huggingface_logo-noborder.sv…"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "huggingface_hub.notebook_login()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Building generator for model `stylegan_animeface512` ...\n",
+      "Finish building generator.\n",
+      "Loading checkpoint from `checkpoints/stylegan_animeface512.pth` ...\n",
+      "Finish loading checkpoint.\n"
+     ]
+    }
+   ],
+   "source": [
+    "animeface_model = utils.load_generator('stylegan_animeface512')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Cloning https://huggingface.co/johko/stylegan_animeface512 into local empty directory.\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "6e51c5ae4a504617aa0f1c1ac798ed15",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Upload file pytorch_model.bin:   0%|          | 32.0k/103M [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "To https://huggingface.co/johko/stylegan_animeface512\n",
+      "   750cd03..2841156  main -> main\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "'https://huggingface.co/johko/stylegan_animeface512/commit/2841156bad3c5a5f47f3edbf4a41880ea8fd3ad3'"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "animeface_model.push_to_hub(\"johko/stylegan_animeface512\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Building generator for model `pggan_celebahq1024` ...\n",
+      "Finish building generator.\n",
+      "Loading checkpoint from `checkpoints/pggan_celebahq1024.pth` ...\n",
+      "Finish loading checkpoint.\n"
+     ]
+    }
+   ],
+   "source": [
+    "celebhq_model = utils.load_generator(\"pggan_celebahq1024\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Cloning https://huggingface.co/johko/pggan-celebahq-1024 into local empty directory.\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "ef4086b23a654b079bd6a3678140c50d",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Upload file pytorch_model.bin:   0%|          | 32.0k/88.1M [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "To https://huggingface.co/johko/pggan-celebahq-1024\n",
+      "   780695e..278449f  main -> main\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "'https://huggingface.co/johko/pggan-celebahq-1024/commit/278449f8416d38a0233c980774528d32c4eee99c'"
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "celebhq_model.push_to_hub(\"johko/pggan-celebahq-1024\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Building generator for model `stylegan_car512` ...\n",
+      "Finish building generator.\n",
+      "Loading checkpoint from `checkpoints/stylegan_car512.pth` ...\n",
+      "Finish loading checkpoint.\n"
+     ]
+    }
+   ],
+   "source": [
+    "cars_model = utils.load_generator(\"stylegan_car512\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 21,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Building generator for model `stylegan_cat256` ...\n",
+      "Finish building generator.\n",
+      "Loading checkpoint from `checkpoints/stylegan_cat256.pth` ...\n",
+      "Finish loading checkpoint.\n"
+     ]
+    }
+   ],
+   "source": [
+    "cats_model = utils.load_generator(\"stylegan_cat256\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Cloning https://huggingface.co/johko/stylegan_cat256 into local empty directory.\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "651e9bff9c9f4555814171195e36d4d3",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Upload file pytorch_model.bin:   0%|          | 32.0k/100M [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "cats_model.push_to_hub(\"johko/stylegan_cat256\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "interpreter": {
+   "hash": "a8d699d01f596cc27ac2722fbc0550b939d217978c7e1ca888dca7ba146ee4bf"
+  },
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

models/model_zoo.py

@@ -9,6 +9,7 @@ MODEL_ZOO = {
         gan_type='pggan',
         resolution=1024,
         url='https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/EW_3jQ6E7xlKvCSHYrbmkQQBAB8tgIv5W5evdT6-GuXiWw?e=gRifVa&download=1',
+        hf_hub_repo='huggan/pggan-celebahq-1024'
     ),
     'pggan_bedroom256': dict(
         gan_type='pggan',
@@ -181,11 +182,13 @@ MODEL_ZOO = {
         gan_type='stylegan',
         resolution=256,
         url='https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/EVjX8u9HuehLip3z0hRfIHcB7QtoFkTB7NiRDb8nrKOl2w?e=lHcp1B&download=1',
+        hf_hub_repo="huggan/stylegan_cat256"
     ),
     'stylegan_car512': dict(
         gan_type='stylegan',
         resolution=512,
         url='https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/EcRJNNzzUzJGjI2X53S9HjkBhXkKT5JRd6Q3IIhCY1AyRw?e=FvMRNj&download=1',
+        hf_hub_repo="huggan/stylegan_car512"
     ),
 
     # StyleGAN ours.
@@ -260,6 +263,7 @@ MODEL_ZOO = {
         gan_type='stylegan',
         resolution=512,
         url='https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/EWDWflY6lBpGgX0CGQpd2Z4B5wTEVamTOA9JRYne7zdCvA?e=tOzgYA&download=1',
+        hf_hub_repo='huggan/stylegan_animeface512'
     ),
     'stylegan_animeportrait512': dict(
         gan_type='stylegan',
@@ -296,15 +300,8 @@ MODEL_ZOO = {
     'stylegan2_car512': dict(
         gan_type='stylegan2',
         resolution=512,
-        url='https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/EYSnUsxU8KJFuMHhZm-JLWoB0nHxdlbrLHNZ_Qkoe3b9LA?e=Ycjp5A&download=1',
+        url='https://mycuhk-my.sharepoint.com/:u:/g/personal/1155082926_link_cuhk_edu_hk/EYSnUsxU8KJFuMHhZm-JLWoB0nHxdlbrLHNZ_Qkoe3b9LA?e=Ycjp5A&download=1'
     ),
-
-    #huggingface models
-    'akhaliq/OneshotCLIP-stylegan2-ffhq' : dict(
-        gan_type='stylegan2',
-        resolution=512,
-        url='akhaliq/OneshotCLIP-stylegan2-ffhq',
-    )
 }
 
 # pylint: enable=line-too-long

models/pggan_generator.py

@@ -6,6 +6,7 @@ Paper: https://arxiv.org/pdf/1710.10196.pdf
 Official TensorFlow implementation:
 https://github.com/tkarras/progressive_growing_of_gans
 """
+import os
 
 import numpy as np
 
@@ -13,6 +14,8 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+from huggingface_hub import PyTorchModelHubMixin, PYTORCH_WEIGHTS_NAME, hf_hub_download
+
 __all__ = ['PGGANGenerator']
 
 # Resolutions allowed.
@@ -25,7 +28,7 @@ _INIT_RES = 4
 _WSCALE_GAIN = np.sqrt(2.0)
 
 
-class PGGANGenerator(nn.Module):
+class PGGANGenerator(nn.Module, PyTorchModelHubMixin):
     """Defines the generator network in PGGAN.
 
     NOTE: The synthesized images are with `RGB` channel order and pixel range
@@ -57,7 +60,8 @@ class PGGANGenerator(nn.Module):
                  fused_scale=False,
                  use_wscale=True,
                  fmaps_base=16 << 10,
-                 fmaps_max=512):
+                 fmaps_max=512,
+                 **kwargs):
         """Initializes with basic settings.
 
         Raises:
@@ -81,6 +85,8 @@ class PGGANGenerator(nn.Module):
         self.use_wscale = use_wscale
         self.fmaps_base = fmaps_base
         self.fmaps_max = fmaps_max
+        
+        self.config = kwargs.pop("config", None)
 
         # Number of convolutional layers.
         self.num_layers = (self.final_res_log2 - self.init_res_log2 + 1) * 2
@@ -202,6 +208,46 @@ class PGGANGenerator(nn.Module):
         }
         return results
 
+    @classmethod
+    def _from_pretrained(
+            cls,
+            model_id,
+            revision,
+            cache_dir,
+            force_download,
+            proxies,
+            resume_download,
+            local_files_only,
+            use_auth_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,
+                use_auth_token=use_auth_token,
+                local_files_only=local_files_only,
+            )
+
+        pretrained = torch.load(model_file, map_location=map_location)
+        return pretrained
+
 
 class PixelNormLayer(nn.Module):
     """Implements pixel-wise feature vector normalization layer."""

models/stylegan2_generator.py

@@ -9,12 +9,14 @@ Paper: https://arxiv.org/pdf/1912.04958.pdf
 
 Official TensorFlow implementation: https://github.com/NVlabs/stylegan2
 """
+import os
 
 import numpy as np
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from huggingface_hub import PYTORCH_WEIGHTS_NAME, hf_hub_download, PyTorchModelHubMixin
 
 from .sync_op import all_gather
 
@@ -33,7 +35,7 @@ _ARCHITECTURES_ALLOWED = ['resnet', 'skip', 'origin']
 _WSCALE_GAIN = 1.0
 
 
-class StyleGAN2Generator(nn.Module):
+class StyleGAN2Generator(nn.Module, PyTorchModelHubMixin):
     """Defines the generator network in StyleGAN2.
 
     NOTE: The synthesized images are with `RGB` channel order and pixel range
@@ -88,7 +90,8 @@ class StyleGAN2Generator(nn.Module):
                  demodulate=True,
                  use_wscale=True,
                  fmaps_base=32 << 10,
-                 fmaps_max=512):
+                 fmaps_max=512,
+                 **kwargs):
         """Initializes with basic settings.
 
         Raises:
@@ -195,6 +198,45 @@ class StyleGAN2Generator(nn.Module):
 
         return {**mapping_results, **synthesis_results}
 
+    @classmethod
+    def _from_pretrained(
+            cls,
+            model_id,
+            revision,
+            cache_dir,
+            force_download,
+            proxies,
+            resume_download,
+            local_files_only,
+            use_auth_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="stylegan2-ffhq-config-f.pt",
+                revision=revision,
+                cache_dir=cache_dir,
+                force_download=force_download,
+                proxies=proxies,
+                resume_download=resume_download,
+                use_auth_token=use_auth_token,
+                local_files_only=local_files_only,
+            )
+
+        pretrained = torch.load(model_file, map_location=map_location)
+        return pretrained
 
 class MappingModule(nn.Module):
     """Implements the latent space mapping module.

models/stylegan_generator.py

@@ -5,6 +5,7 @@ Paper: https://arxiv.org/pdf/1812.04948.pdf
 
 Official TensorFlow implementation: https://github.com/NVlabs/stylegan
 """
+import os
 
 import numpy as np
 
@@ -14,6 +15,8 @@ import torch.nn.functional as F
 
 from .sync_op import all_gather
 
+from huggingface_hub import PyTorchModelHubMixin, PYTORCH_WEIGHTS_NAME, hf_hub_download
+
 __all__ = ['StyleGANGenerator']
 
 # Resolutions allowed.
@@ -33,7 +36,7 @@ _WSCALE_GAIN = np.sqrt(2.0)
 _STYLEMOD_WSCALE_GAIN = 1.0
 
 
-class StyleGANGenerator(nn.Module):
+class StyleGANGenerator(nn.Module, PyTorchModelHubMixin):
     """Defines the generator network in StyleGAN.
 
     NOTE: The synthesized images are with `RGB` channel order and pixel range
@@ -83,7 +86,8 @@ class StyleGANGenerator(nn.Module):
                  fused_scale='auto',
                  use_wscale=True,
                  fmaps_base=16 << 10,
-                 fmaps_max=512):
+                 fmaps_max=512,
+                 **kwargs):
         """Initializes with basic settings.
 
         Raises:
@@ -115,6 +119,9 @@ class StyleGANGenerator(nn.Module):
         self.use_wscale = use_wscale
         self.fmaps_base = fmaps_base
         self.fmaps_max = fmaps_max
+        
+        self.config = kwargs.pop("config", None)
+
 
         self.num_layers = int(np.log2(self.resolution // self.init_res * 2)) * 2
 
@@ -188,6 +195,46 @@ class StyleGANGenerator(nn.Module):
 
         return {**mapping_results, **synthesis_results}
 
+    @classmethod
+    def _from_pretrained(
+            cls,
+            model_id,
+            revision,
+            cache_dir,
+            force_download,
+            proxies,
+            resume_download,
+            local_files_only,
+            use_auth_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,
+                use_auth_token=use_auth_token,
+                local_files_only=local_files_only,
+            )
+
+        pretrained = torch.load(model_file, map_location=map_location)
+        return pretrained
+
 
 class MappingModule(nn.Module):
     """Implements the latent space mapping module.

utils.py

@@ -50,7 +50,7 @@ def postprocess(images, min_val=-1.0, max_val=1.0):
     return images
 
 
-def load_generator(model_name):
+def load_generator(model_name, from_hf_hub=False):
     """Loads pre-trained generator.
 
     Args:
@@ -74,19 +74,25 @@ def load_generator(model_name):
     generator = build_generator(**model_config)
     print(f'Finish building generator.')
 
-    # Load pre-trained weights.
-    os.makedirs(CHECKPOINT_DIR, exist_ok=True)
-    checkpoint_path = os.path.join(CHECKPOINT_DIR, model_name + '.pth')
-    print(f'Loading checkpoint from `{checkpoint_path}` ...')
-    if not os.path.exists(checkpoint_path):
-        print(f'  Downloading checkpoint from `{url}` ...')
-        subprocess.call(['wget', '--quiet', '-O', checkpoint_path, url])
-        print(f'  Finish downloading checkpoint.')
-    checkpoint = torch.load(checkpoint_path, map_location='cpu')
-    if 'generator_smooth' in checkpoint:
-        generator.load_state_dict(checkpoint['generator_smooth'])
+    if from_hf_hub and "hf_hub_repo" in model_config.keys():
+        checkpoint = generator.from_pretrained(model_config["hf_hub_repo"])
+        generator.load_state_dict(checkpoint)
+        print("loaded from hf_hub")
     else:
-        generator.load_state_dict(checkpoint['generator'])
+        # Load pre-trained weights.
+        os.makedirs(CHECKPOINT_DIR, exist_ok=True)
+        checkpoint_path = os.path.join(CHECKPOINT_DIR, model_name + '.pth')
+        print(f'Loading checkpoint from `{checkpoint_path}` ...')
+        if not os.path.exists(checkpoint_path):
+            print(f'  Downloading checkpoint from `{url}` ...')
+            subprocess.call(['wget', '--quiet', '-O', checkpoint_path, url])
+            print(f'  Finish downloading checkpoint.')
+        checkpoint = torch.load(checkpoint_path, map_location='cpu')
+
+        if 'generator_smooth' in checkpoint:
+            generator.load_state_dict(checkpoint['generator_smooth'])
+        else:
+            generator.load_state_dict(checkpoint['generator'])
     #generator = generator.cuda()
     generator.eval()
     print(f'Finish loading checkpoint.')