Update

.pre-commit-config.yaml

@@ -21,26 +21,17 @@ repos:
   - id: docformatter
     args: ['--in-place']
 - repo: https://github.com/pycqa/isort
-  rev: 5.10.1
+  rev: 5.12.0
   hooks:
     - id: isort
 - repo: https://github.com/pre-commit/mirrors-mypy
-  rev: v0.812
+  rev: v0.991
   hooks:
     - id: mypy
       args: ['--ignore-missing-imports']
+      additional_dependencies: ['types-python-slugify']
 - repo: https://github.com/google/yapf
   rev: v0.32.0
   hooks:
   - id: yapf
     args: ['--parallel', '--in-place']
-- repo: https://github.com/kynan/nbstripout
-  rev: 0.5.0
-  hooks:
-    - id: nbstripout
-      args: ['--extra-keys', 'metadata.interpreter metadata.kernelspec cell.metadata.pycharm']
-- repo: https://github.com/nbQA-dev/nbQA
-  rev: 1.3.1
-  hooks:
-    - id: nbqa-isort
-    - id: nbqa-yapf

README.md

@@ -4,9 +4,12 @@ emoji: 🐨
 colorFrom: indigo
 colorTo: red
 sdk: gradio
-sdk_version: 3.1.3
+sdk_version: 3.35.2
 app_file: app.py
 pinned: false
+suggested_hardware: t4-small
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces#reference
+
+https://arxiv.org/abs/2202.12211

app.py

@@ -2,7 +2,6 @@
 
 from __future__ import annotations
 
-import argparse
 import pathlib
 
 import gradio as gr
@@ -10,23 +9,7 @@ import numpy as np
 
 from model import Model
 
-TITLE = '# Self-Distilled StyleGAN'
-DESCRIPTION = '''This is an unofficial demo for [https://github.com/self-distilled-stylegan/self-distilled-internet-photos](https://github.com/self-distilled-stylegan/self-distilled-internet-photos).
-
-Expected execution time on Hugging Face Spaces: 2s'''
-FOOTER = '<img id="visitor-badge" src="https://visitor-badge.glitch.me/badge?page_id=hysts.self-distilled-stylegan" alt="visitor badge" />'
-
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser()
-    parser.add_argument('--device', type=str, default='cpu')
-    parser.add_argument('--theme', type=str)
-    parser.add_argument('--share', action='store_true')
-    parser.add_argument('--port', type=int)
-    parser.add_argument('--disable-queue',
-                        dest='enable_queue',
-                        action='store_false')
-    return parser.parse_args()
+DESCRIPTION = '# [Self-Distilled StyleGAN](https://github.com/self-distilled-stylegan/self-distilled-internet-photos)'
 
 
 def get_sample_image_url(name: str) -> str:
@@ -56,87 +39,71 @@ def get_cluster_center_image_markdown(model_name: str) -> str:
     return f'![cluster center images]({url})'
 
 
-def main():
-    args = parse_args()
-
-    model = Model(args.device)
-
-    with gr.Blocks(theme=args.theme, css='style.css') as demo:
-        gr.Markdown(TITLE)
-        gr.Markdown(DESCRIPTION)
-
-        with gr.Tabs():
-            with gr.TabItem('App'):
-                with gr.Row():
-                    with gr.Column():
-                        with gr.Group():
-                            model_name = gr.Dropdown(
-                                model.MODEL_NAMES,
-                                value=model.MODEL_NAMES[0],
-                                label='Model')
-                            seed = gr.Slider(0,
-                                             np.iinfo(np.uint32).max,
-                                             value=0,
-                                             step=1,
-                                             label='Seed')
-                            psi = gr.Slider(0,
-                                            2,
-                                            step=0.05,
-                                            value=0.7,
-                                            label='Truncation psi')
-                            truncation_type = gr.Dropdown(
-                                model.TRUNCATION_TYPES,
-                                value=model.TRUNCATION_TYPES[0],
-                                label='Truncation Type')
-                            run_button = gr.Button('Run')
-                    with gr.Column():
-                        result = gr.Image(label='Result', elem_id='result')
-
-            with gr.TabItem('Sample Images'):
-                with gr.Row():
-                    paths = sorted(pathlib.Path('samples').glob('*'))
-                    names = [path.stem for path in paths]
-                    model_name2 = gr.Dropdown(
-                        names,
-                        value='dogs_1024_multimodal_lpips',
-                        label='Type')
-                with gr.Row():
-                    text = get_sample_image_markdown(model_name2.value)
-                    sample_images = gr.Markdown(text)
-
-            with gr.TabItem('Cluster Center Images'):
-                with gr.Row():
-                    model_name3 = gr.Dropdown(model.MODEL_NAMES,
-                                              value=model.MODEL_NAMES[0],
-                                              label='Model')
-                with gr.Row():
-                    text = get_cluster_center_image_markdown(model_name3.value)
-                    cluster_center_images = gr.Markdown(value=text)
-
-        gr.Markdown(FOOTER)
-
-        model_name.change(fn=model.set_model, inputs=model_name, outputs=None)
-        run_button.click(fn=model.set_model_and_generate_image,
-                         inputs=[
-                             model_name,
-                             seed,
-                             psi,
-                             truncation_type,
-                         ],
-                         outputs=result)
-        model_name2.change(fn=get_sample_image_markdown,
-                           inputs=model_name2,
-                           outputs=sample_images)
-        model_name3.change(fn=get_cluster_center_image_markdown,
-                           inputs=model_name3,
-                           outputs=cluster_center_images)
-
-    demo.launch(
-        enable_queue=args.enable_queue,
-        server_port=args.port,
-        share=args.share,
-    )
-
-
-if __name__ == '__main__':
-    main()
+model = Model()
+
+with gr.Blocks(css='style.css') as demo:
+    gr.Markdown(DESCRIPTION)
+
+    with gr.Tabs():
+        with gr.TabItem('App'):
+            with gr.Row():
+                with gr.Column():
+                    with gr.Group():
+                        model_name = gr.Dropdown(label='Model',
+                                                 choices=model.MODEL_NAMES,
+                                                 value=model.MODEL_NAMES[0])
+                        seed = gr.Slider(label='Seed',
+                                         minimum=0,
+                                         maximum=np.iinfo(np.uint32).max,
+                                         step=1,
+                                         value=0)
+                        psi = gr.Slider(label='Truncation psi',
+                                        minimum=0,
+                                        maximum=2,
+                                        step=0.05,
+                                        value=0.7)
+                        truncation_type = gr.Dropdown(
+                            label='Truncation Type',
+                            choices=model.TRUNCATION_TYPES,
+                            value=model.TRUNCATION_TYPES[0])
+                        run_button = gr.Button('Run')
+                with gr.Column():
+                    result = gr.Image(label='Result', elem_id='result')
+
+        with gr.TabItem('Sample Images'):
+            with gr.Row():
+                paths = sorted(pathlib.Path('samples').glob('*'))
+                names = [path.stem for path in paths]
+                model_name2 = gr.Dropdown(label='Type',
+                                          choices=names,
+                                          value='dogs_1024_multimodal_lpips')
+            with gr.Row():
+                text = get_sample_image_markdown(model_name2.value)
+                sample_images = gr.Markdown(text)
+
+        with gr.TabItem('Cluster Center Images'):
+            with gr.Row():
+                model_name3 = gr.Dropdown(label='Model',
+                                          choices=model.MODEL_NAMES,
+                                          value=model.MODEL_NAMES[0])
+            with gr.Row():
+                text = get_cluster_center_image_markdown(model_name3.value)
+                cluster_center_images = gr.Markdown(value=text)
+
+    model_name.change(fn=model.set_model, inputs=model_name)
+    run_button.click(fn=model.set_model_and_generate_image,
+                     inputs=[
+                         model_name,
+                         seed,
+                         psi,
+                         truncation_type,
+                     ],
+                     outputs=result)
+    model_name2.change(fn=get_sample_image_markdown,
+                       inputs=model_name2,
+                       outputs=sample_images)
+    model_name3.change(fn=get_cluster_center_image_markdown,
+                       inputs=model_name3,
+                       outputs=cluster_center_images)
+
+demo.queue(max_size=10).launch()

model.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import os
 import pathlib
 import pickle
 import sys
@@ -15,8 +14,6 @@ current_dir = pathlib.Path(__file__).parent
 submodule_dir = current_dir / 'stylegan3'
 sys.path.insert(0, submodule_dir.as_posix())
 
-HF_TOKEN = os.environ['HF_TOKEN']
-
 
 class LPIPS(lpips.LPIPS):
     @staticmethod
@@ -60,8 +57,9 @@ class Model:
         'Global',
     ]
 
-    def __init__(self, device: str | torch.device):
-        self.device = torch.device(device)
+    def __init__(self):
+        self.device = torch.device(
+            'cuda:0' if torch.cuda.is_available() else 'cpu')
         self._download_all_models()
         self._download_all_cluster_centers()
         self._download_all_cluster_center_images()
@@ -77,9 +75,8 @@ class Model:
         )
 
     def _load_model(self, model_name: str) -> nn.Module:
-        path = hf_hub_download('hysts/Self-Distilled-StyleGAN',
-                               f'models/{model_name}_pytorch.pkl',
-                               use_auth_token=HF_TOKEN)
+        path = hf_hub_download('public-data/Self-Distilled-StyleGAN',
+                               f'models/{model_name}_pytorch.pkl')
         with open(path, 'rb') as f:
             model = pickle.load(f)['G_ema']
         model.eval()
@@ -87,17 +84,15 @@ class Model:
         return model
 
     def _load_cluster_centers(self, model_name: str) -> torch.Tensor:
-        path = hf_hub_download('hysts/Self-Distilled-StyleGAN',
-                               f'cluster_centers/{model_name}.npy',
-                               use_auth_token=HF_TOKEN)
+        path = hf_hub_download('public-data/Self-Distilled-StyleGAN',
+                               f'cluster_centers/{model_name}.npy')
         centers = np.load(path)
         centers = torch.from_numpy(centers).float().to(self.device)
         return centers
 
     def _load_cluster_center_images(self, model_name: str) -> np.ndarray:
-        path = hf_hub_download('hysts/Self-Distilled-StyleGAN',
-                               f'cluster_center_images/{model_name}.npy',
-                               use_auth_token=HF_TOKEN)
+        path = hf_hub_download('public-data/Self-Distilled-StyleGAN',
+                               f'cluster_center_images/{model_name}.npy')
         return np.load(path)
 
     def set_model(self, model_name: str) -> None:

requirements.txt

@@ -1,6 +1,6 @@
 lpips==0.1.4
-numpy==1.22.3
-Pillow==9.0.1
-scipy==1.8.0
-torch==1.11.0
-torchvision==0.12.0
+numpy==1.23.5
+Pillow==10.0.0
+scipy==1.10.1
+torch==2.0.1
+torchvision==0.15.2