Update app.py

app.py

@@ -59,8 +59,8 @@ def generate_image(seed: int, truncation_psi: float, model: nn.Module,
     return out[0].cpu().numpy()
 
 
-def load_model(file_name: str, device: torch.device) -> nn.Module:
-    path = hf_hub_download('feng2022/Time-TravelRephotography',
+def load_model(path:str, file_name: str, device: torch.device) -> nn.Module:
+    path = hf_hub_download(f'{path}',
                            f'{file_name}',
                            use_auth_token=TOKEN)
     with open(path, 'rb') as f:
@@ -78,7 +78,7 @@ def main():
     args = parse_args()
     device = torch.device(args.device)
 
-    model = load_model('stylegan_human_v2_1024.pkl', device)
+    model = load_model('feng2022/Time-TravelRephotography_e4e_ffhq_encode','e4e_ffhq_encode.pt', device)
 
     func = functools.partial(generate_image, model=model, device=device)
     func = functools.update_wrapper(func, generate_image)