Update

.pre-commit-config.yaml

@@ -0,0 +1,46 @@
+exclude: ^DualStyleGAN
+repos:
+- repo: https://github.com/pre-commit/pre-commit-hooks
+  rev: v4.2.0
+  hooks:
+  - id: check-executables-have-shebangs
+  - id: check-json
+  - id: check-merge-conflict
+  - id: check-shebang-scripts-are-executable
+  - id: check-toml
+  - id: check-yaml
+  - id: double-quote-string-fixer
+  - id: end-of-file-fixer
+  - id: mixed-line-ending
+    args: ['--fix=lf']
+  - id: requirements-txt-fixer
+  - id: trailing-whitespace
+- repo: https://github.com/myint/docformatter
+  rev: v1.4
+  hooks:
+  - id: docformatter
+    args: ['--in-place']
+- repo: https://github.com/pycqa/isort
+  rev: 5.10.1
+  hooks:
+    - id: isort
+- repo: https://github.com/pre-commit/mirrors-mypy
+  rev: v0.812
+  hooks:
+    - id: mypy
+      args: ['--ignore-missing-imports']
+- 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

.style.yapf

@@ -0,0 +1,5 @@
+[style]
+based_on_style = pep8
+blank_line_before_nested_class_or_def = false
+spaces_before_comment = 2
+split_before_logical_operator = true

app.py

@@ -9,6 +9,12 @@ import gradio as gr
 
 from dualstylegan import Model
 
+DESCRIPTION = '''# Portrait Style Transfer with <a href="https://github.com/williamyang1991/DualStyleGAN">DualStyleGAN</a>
+
+<img id="overview" alt="overview" src="https://raw.githubusercontent.com/williamyang1991/DualStyleGAN/main/doc_images/overview.jpg" />
+'''
+FOOTER = '<img id="visitor-badge" alt="visitor badge" src="https://visitor-badge.glitch.me/badge?page_id=gradio-blocks.dualstylegan" />'
+
 
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser()
@@ -81,26 +87,9 @@ def main():
     args = parse_args()
     model = Model(device=args.device)
 
-    css = '''
-h1#title {
-  text-align: center;
-}
-img#overview {
-  max-width: 800px;
-  max-height: 600px;
-}
-img#style-image {
-  max-width: 1000px;
-  max-height: 600px;
-}
-'''
-
-    with gr.Blocks(theme=args.theme, css=css) as demo:
-        gr.Markdown(
-            '''<h1 id="title">Portrait Style Transfer with <a href="https://github.com/williamyang1991/DualStyleGAN">DualStyleGAN</a></h1>
+    with gr.Blocks(theme=args.theme, css='style.css') as demo:
+        gr.Markdown(DESCRIPTION)
 
-<center><img id="overview" src="https://raw.githubusercontent.com/williamyang1991/DualStyleGAN/main/doc_images/overview.jpg" alt="overview"></center>
-''')
         with gr.Box():
             gr.Markdown('''## Step 1 (Preprocess Input Image)
 
@@ -202,9 +191,7 @@ img#style-image {
                         [1.0, 0.0],
                     ])
 
-        gr.Markdown(
-            '<center><img src="https://visitor-badge.glitch.me/badge?page_id=gradio-blocks.dualstylegan" alt="visitor badge"/></center>'
-        )
+        gr.Markdown(FOOTER)
 
         detect_button.click(fn=model.detect_and_align_face,
                             inputs=input_image,

dualstylegan.py

@@ -23,12 +23,11 @@ from model.dualstylegan import DualStyleGAN
 from model.encoder.align_all_parallel import align_face
 from model.encoder.psp import pSp
 
-TOKEN = os.environ['TOKEN']
+HF_TOKEN = os.environ['HF_TOKEN']
 MODEL_REPO = 'hysts/DualStyleGAN'
 
 
 class Model:
-
     def __init__(self, device: Union[torch.device, str]):
         self.device = torch.device(device)
         self.landmark_model = self._create_dlib_landmark_model()
@@ -58,13 +57,13 @@ class Model:
         path = huggingface_hub.hf_hub_download(
             'hysts/dlib_face_landmark_model',
             'shape_predictor_68_face_landmarks.dat',
-            use_auth_token=TOKEN)
+            use_auth_token=HF_TOKEN)
         return dlib.shape_predictor(path)
 
     def _load_encoder(self) -> nn.Module:
         ckpt_path = huggingface_hub.hf_hub_download(MODEL_REPO,
                                                     'models/encoder.pt',
-                                                    use_auth_token=TOKEN)
+                                                    use_auth_token=HF_TOKEN)
         ckpt = torch.load(ckpt_path, map_location='cpu')
         opts = ckpt['opts']
         opts['device'] = self.device.type
@@ -90,7 +89,7 @@ class Model:
         ckpt_path = huggingface_hub.hf_hub_download(
             MODEL_REPO,
             f'models/{style_type}/generator.pt',
-            use_auth_token=TOKEN)
+            use_auth_token=HF_TOKEN)
         ckpt = torch.load(ckpt_path, map_location='cpu')
         model.load_state_dict(ckpt['g_ema'])
         model.to(self.device)
@@ -106,7 +105,7 @@ class Model:
         path = huggingface_hub.hf_hub_download(
             MODEL_REPO,
             f'models/{style_type}/{filename}',
-            use_auth_token=TOKEN)
+            use_auth_token=HF_TOKEN)
         exstyles = np.load(path, allow_pickle=True).item()
         return exstyles
 

style.css

@@ -0,0 +1,17 @@
+h1 {
+  text-align: center;
+}
+img#overview {
+  max-width: 800px;
+  max-height: 600px;
+  display: block;
+  margin: auto;
+}
+img#style-image {
+  max-width: 1000px;
+  max-height: 600px;
+}
+img#visitor-badge {
+  display: block;
+  margin: auto;
+}