Update src/pixel3dmm/preprocessing/MICA/demo.py

src/pixel3dmm/preprocessing/MICA/demo.py

@@ -1,3 +1,6 @@
+
+
+
 # -*- coding: utf-8 -*-
 
 # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
@@ -131,18 +134,11 @@ def load_checkpoint(args, mica):
         mica.flameModel.load_state_dict(checkpoint['flameModel'])
 
 
-def main(cfg, args):
-    device = 'cuda:0'
-    cfg.model.testing = True
-    mica = util.find_model_using_name(model_dir='micalib.models', model_name=cfg.model.name)(cfg, device)
-    load_checkpoint(args, mica)
-    mica.eval()
+def main(args, mica, app):
 
     faces = mica.flameModel.generator.faces_tensor.cpu()
     Path(args.o).mkdir(exist_ok=True, parents=True)
 
-    app = LandmarksDetector(model=detectors.RETINAFACE)
-
     with torch.no_grad():
         logger.info(f'Processing has started...')
         paths = process(args, app, draw_bbox=False)
@@ -186,4 +182,14 @@ if __name__ == '__main__':
                             <<<<<<<< ALREADY COMPLETE MICA PREDICTION FOR {args.video_name}, SKIPPING >>>>>>>>
                             ''')
             exit()
-    main(cfg, args)
+
+    # instantiate models outside main
+    device = 'cuda'
+    cfg.model.testing = True
+    mica = util.find_model_using_name(model_dir='micalib.models', model_name=cfg.model.name)(cfg, device)
+    load_checkpoint(args, mica)
+    mica.eval()
+
+    app = LandmarksDetector(model=detectors.RETINAFACE)
+
+    main(args, mica, app)