Update app.py

app.py

@@ -54,15 +54,11 @@ def gen_sources(deepfake_img):
     #------------Initialize:Decoder-F------------------------
     config_path = "./models/config.yaml"
     checkpoint_path_f = "./models/model_vaq1_ff.pth"
-    checkpoint_f = torch.load(checkpoint_path_f, map_location=device)
-    model_vaq_f = Generator(config_path, device)
-    model_vaq_f = model_vaq_f.load_state_dict(checkpoint_f, strict=True)
+    model_vaq_f = Generator(config_path, checkpoint_path_f, device)
     model_vaq_f.eval()
     #------------Initialize:Decoder-G------------------------
     checkpoint_path_g = "./models/model_vaq2_gg.pth"
-    checkpoint_g = torch.load(checkpoint_path_g, map_location=device)
-    model_vaq_g = Generator(config_path, device)
-    model_vaq_g = model_vaq_g.load_state_dict(checkpoint_g, strict=True)
+    model_vaq_g = Generator(config_path, checkpoint_path_g, device)
     model_vaq_g.eval()
     ##------------------------Initialize Model-F-------------------------------------
     model_z1 = DeepfakeToSourceTransformer().to(device)