Update app.py

app.py

@@ -9,6 +9,7 @@ import stylegan2
 from stylegan2 import utils
 from huggingface_hub import hf_hub_download
 import gradio as gr
+import re
 
 from types import SimpleNamespace
 
@@ -78,9 +79,17 @@ def generate_images(G, args):
  
  
  #----------------------------------------------------------------------------
+interface_modelversion_labels = [
+    "TWDNEv3 iteration 24664 (best and current version on TWDNE)",
+    "TWDNEv3 iteration 18528 (the most used version on the Internet)",
+    "TWDNEv3 iteration 17325"
+]
 
-def inference(seed, truncation_psi, model_version):
-    G = stylegan2.models.load(hf_hub_download("hr16/Gwern-TWDNEv3-pytorch_ckpt", f"{model_version}/Gs.pth", use_auth_token=os.environ['MODEL_READING_TOKEN']))
+def inference(seed, truncation_psi, modelversion_label):
+    model_iteration = re.search("TWDNEv3 iteration (\d{5})", modelversion_label).group(1)
+    G = stylegan2.models.load(
+        hf_hub_download("hr16/Gwern-TWDNEv3-pytorch_ckpt", f"iteration-{model_iteration}/Gs.pth", use_auth_token=os.environ['MODEL_READING_TOKEN'])
+    )
     G.eval()
     return generate_images(
         G, 
@@ -103,14 +112,10 @@ gr.Interface(
         gr.Number(precision=0, label="PCG64 PRNG Seed (any-bit-size unsigned int, note that it may different from the original site)"), 
         gr.Slider(0, 2, step=0.1, value=0.7, label='Truncation psi (aka creative level, between 0 and 2)'),
         gr.Radio(
-            ["iteration-24664", "iteration-18528", "iteration-17325"], 
-            value="iteration-24664", 
+            interface_modelversion_labels,
+            value="TWDNEv3 iteration 24664 (best and current version on TWDNE)", 
             type="value", 
-            label=[
-                "TWDNEv3 iteration 24664 (best and current version on TWDNE)", 
-                "TWDNEv3 iteration 18528 (the most used version on Internet)", 
-                "TWDNEv3 iteration 17325"
-            ]           
+            label="Model versions"         
         )
     ], 
     gr.outputs.Image(type="pil"),