Updated to use arbitrary model paths

app.py

@@ -31,25 +31,22 @@ from generate_videos import generate_frames, video_from_interpolations, vid_to_g
 model_dir = "models"
 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": "stylegan2-ffhq-config-f.pt"}
+model_repos = {"e4e": ("akhaliq/JoJoGAN_e4e_ffhq_encode", "e4e_ffhq_encode.pt"),
+               "dlib": ("akhaliq/jojogan_dlib", "shape_predictor_68_face_landmarks.dat"),
+               "base": ("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_paths = {}
 
-    model_list = [Path(model_ckpt).stem for model_ckpt in os.listdir(model_dir)]
+    for model_name, repo_details in model_repos.items():
+        download_path = hf_hub_download(repo_id=repo_details[0], filename=repo_details[1])
+        model_paths[model_name] = download_path
 
-    return model_list
+    return model_paths
 
-model_list = get_models()
+model_paths = get_models()
 
 class ImageEditor(object):
     def __init__(self):
@@ -62,18 +59,20 @@ class ImageEditor(object):
 
         self.generators = {}
 
-        for model in model_list:
+        self.model_list = [name for name in model_paths.keys() if name not in ["e4e", "dlib"]]
+
+        for model in self.model_list:
             g_ema = Generator(
                 model_size, latent_size, n_mlp, channel_multiplier=channel_mult
             ).to(self.device)
 
-            checkpoint = torch.load(f"models/{model}.pt")
+            checkpoint = torch.load(model_paths[model])
 
             g_ema.load_state_dict(checkpoint['g_ema'])
 
             self.generators[model] = g_ema
 
-        self.experiment_args = {"model_path": "e4e_ffhq_encode.pt"}
+        self.experiment_args = {"model_path": model_paths["e4e"]}
         self.experiment_args["transform"] = transforms.Compose(
             [
                 transforms.Resize((256, 256)),
@@ -96,7 +95,7 @@ class ImageEditor(object):
         self.e4e_net.cuda()
 
         self.shape_predictor = dlib.shape_predictor(
-            models_and_paths["akhaliq/jojogan_dlib"]
+            model_paths["dlib"]
         )
 
         print("setup complete")
@@ -120,11 +119,11 @@ class ImageEditor(object):
     ):  
 
         if output_style == 'all':
-            styles = model_list
+            styles = self.model_list
         elif output_style == 'list - enter below':
             styles = style_list.split(",")
             for style in styles:
-                if style not in model_list:
+                if style not in self.model_list:
                     raise ValueError(f"Encountered style '{style}' in the style_list which is not an available option.")
         else:
             styles = [output_style]