Update

app.py

@@ -26,23 +26,6 @@ from model.dualstylegan import DualStyleGAN
 from model.encoder.align_all_parallel import align_face
 from model.encoder.psp import pSp
 
-STYLE_IMAGE_PATHS = {
-    'cartoon':
-    'https://raw.githubusercontent.com/williamyang1991/DualStyleGAN/main/doc_images/cartoon_overview.jpg',
-    'caricature':
-    'https://raw.githubusercontent.com/williamyang1991/DualStyleGAN/main/doc_images/caricature_overview.jpg',
-    'anime':
-    'https://raw.githubusercontent.com/williamyang1991/DualStyleGAN/main/doc_images/anime_overview.jpg',
-    'arcane':
-    'https://raw.githubusercontent.com/williamyang1991/DualStyleGAN/main/doc_images/Reconstruction_arcane_overview.jpg',
-    'comic':
-    'https://raw.githubusercontent.com/williamyang1991/DualStyleGAN/main/doc_images/Reconstruction_comic_overview.jpg',
-    'pixar':
-    'https://raw.githubusercontent.com/williamyang1991/DualStyleGAN/main/doc_images/Reconstruction_pixar_overview.jpg',
-    'slamdunk':
-    'https://raw.githubusercontent.com/williamyang1991/DualStyleGAN/main/doc_images/Reconstruction_slamdunk_overview.jpg',
-}
-
 TOKEN = os.environ['TOKEN']
 MODEL_REPO = 'hysts/DualStyleGAN'
 
@@ -199,6 +182,25 @@ class App:
         return img_gen
 
 
+def get_style_image_url(style_name: str) -> str:
+    base_url = 'https://raw.githubusercontent.com/williamyang1991/DualStyleGAN/main/doc_images'
+    filenames = {
+        'cartoon': 'cartoon_overview.jpg',
+        'caricature': 'caricature_overview.jpg',
+        'anime': 'anime_overview.jpg',
+        'arcane': 'Reconstruction_arcane_overview.jpg',
+        'comic': 'Reconstruction_comic_overview.jpg',
+        'pixar': 'Reconstruction_pixar_overview.jpg',
+        'slamdunk': 'Reconstruction_slamdunk_overview.jpg',
+    }
+    return f'{base_url}/{filenames[style_name]}'
+
+
+def get_style_image_markdown_text(style_name: str) -> str:
+    url = get_style_image_url(style_name)
+    return f'<center><img id="style-image" src="{url}" alt="style image"></center>'
+
+
 def update_slider(choice: str) -> dict:
     max_vals = {
         'cartoon': 316,
@@ -212,9 +214,8 @@ def update_slider(choice: str) -> dict:
     return gr.Slider.update(maximum=max_vals[choice] + 1, value=26)
 
 
-def update_style_image(choice: str) -> dict:
-    style_image_path = STYLE_IMAGE_PATHS[choice]
-    text = f'<center><img src="{style_image_path}" alt="style image" width="800" height="400"></center>'
+def update_style_image(style_name: str) -> dict:
+    text = get_style_image_markdown_text(style_name)
     return gr.Markdown.update(value=text)
 
 
@@ -222,13 +223,27 @@ def main():
     args = parse_args()
     app = App(device=torch.device(args.device))
 
-    with gr.Blocks(theme=args.theme) as demo:
+    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(
-            '''<center><h1>Portrait Style Transfer with DualStyleGAN</h1></center>
+            '''<h1 id="title">Portrait Style Transfer with DualStyleGAN</h1>
 
 This is an unofficial demo app for https://github.com/williamyang1991/DualStyleGAN.
 
-<center><img src="https://raw.githubusercontent.com/williamyang1991/DualStyleGAN/main/doc_images/overview.jpg" alt="overview" width="800" height="400"></center>
+<center><img id="overview" src="https://raw.githubusercontent.com/williamyang1991/DualStyleGAN/main/doc_images/overview.jpg" alt="overview"></center>
 
 Related App: https://huggingface.co/spaces/hysts/DualStyleGAN
 ''')
@@ -278,8 +293,7 @@ Related App: https://huggingface.co/spaces/hysts/DualStyleGAN
                                                 step=1,
                                                 label='Style Image Index',
                                                 interactive=True)
-                    style_image_path = STYLE_IMAGE_PATHS['cartoon']
-                    text = f'<center><img src="{style_image_path}" alt="style image" width="800" height="400"></center>'
+                    text = get_style_image_markdown_text('cartoon')
                     style_image = gr.Markdown(value=text)
 
         with gr.Box():