Fixed model downloads

app.py

@@ -17,6 +17,7 @@ import tempfile
 import shutil
 from argparse import Namespace
 from pathlib import Path
+import shutil
 
 import dlib
 import numpy as np
@@ -32,13 +33,19 @@ os.makedirs(model_dir, exist_ok=True)
 
 models_and_paths = {"akhaliq/JoJoGAN_e4e_ffhq_encode": "e4e_ffhq_encode.pt",
                     "akhaliq/jojogan_dlib": "shape_predictor_68_face_landmarks.dat",
-                    "akhaliq/jojogan-stylegan2-ffhq-config-f": f"{model_dir}/base.pt"}
+                    "akhaliq/jojogan-stylegan2-ffhq-config-f": "stylegan2-ffhq-config-f.pt"}
 
 def get_models():
+    os.makedirs(model_dir, exist_ok=True)
+
     for repo_id, file_path in models_and_paths.items():
         hf_hub_download(repo_id=repo_id, filename=file_path)
+        if not "akhaliq" in repo_id:
+            shutil.move(file_path, os.path.join(model_dir, file_path))
+        elif "stylegan2" in file_path:
+            shutil.move(file_path, os.path.join(model_dir, "base.pt"))    
 
-    model_list = ['base'] + [Path(model_ckpt).stem for model_ckpt in os.listdir(model_dir) if not 'base' in model_ckpt]
+    model_list = [Path(model_ckpt).stem for model_ckpt in os.listdir(model_dir)]
 
     return model_list