Update app.py

app.py

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