init project

app.py

@@ -10,6 +10,7 @@ import argparse
 
 from modules.pe3r.demo import main_demo
 from modules.pe3r.models import Models
+import torch
 
 # def set_print_with_timestamp(time_format="%Y-%m-%d %H:%M:%S"):
 #     builtin_print = builtins.print
@@ -20,6 +21,8 @@ from modules.pe3r.models import Models
 #         builtin_print(*args, **kwargs)
 #     builtins.print = print_with_timestamp
 
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
 def get_args_parser():
     parser = argparse.ArgumentParser()
     parser_url = parser.add_mutually_exclusive_group()
@@ -28,7 +31,7 @@ def get_args_parser():
     parser_url.add_argument("--server_name", type=str, default=None, help="server url, default is 127.0.0.1")
     parser.add_argument("--server_port", type=int, help=("will start gradio app on this port (if available). "
                                                          "If None, will search for an available port starting at 7860."), default=None)
-    parser.add_argument("--device", type=str, default='cuda', help="pytorch device")
+    # parser.add_argument("--device", type=str, default='cuda', help="pytorch device")
     parser.add_argument("--tmp_dir", type=str, default=None, help="value for tempfile.tempdir")
     parser.add_argument("--silent", action='store_true', default=False, help="silence logs")
     # change defaults
@@ -50,9 +53,9 @@ if __name__ == '__main__':
     else:
         server_name = '0.0.0.0' if args.local_network else '127.0.0.1'
 
-    pe3r = Models(device=args.device)
+    pe3r = Models(device=device)
 
     with tempfile.TemporaryDirectory(suffix='pe3r_gradio_demo') as tmpdirname:
         if not args.silent:
             print('Outputing stuff in', tmpdirname)
-        main_demo(tmpdirname, pe3r, args.device, server_name, args.server_port, silent=args.silent)
+        main_demo(tmpdirname, pe3r, device, server_name, args.server_port, silent=args.silent)