update

app.py

@@ -16,6 +16,7 @@ import sys
 import os
 import pathlib
 
+import argparse
 import gradio as gr
 import torch
 
@@ -24,6 +25,19 @@ from inference import inference_fn
 # from trainer import Trainer
 # from uploader import upload
 
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--device', type=str, default='cpu')
+    parser.add_argument('--theme', type=str)
+    parser.add_argument('--share', action='store_true')
+    parser.add_argument('--port', type=int)
+    parser.add_argument('--disable-queue',
+                        dest='enable_queue',
+                        action='store_false')
+    return parser.parse_args()
+
+
 TITLE = '# ReVersion'
 DESCRIPTION = '''This is a demo for [https://github.com/ziqihuangg/ReVersion](https://github.com/ziqihuangg/ReVersion).
 It is recommended to upgrade to GPU in Settings after duplicating this space to use it.
@@ -149,6 +163,10 @@ def create_inference_demo(func: inference_fn) -> gr.Blocks:
     return demo
 
 
+args = parse_args()
+args.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+print('*** Now using %s.'%(args.device))
+
 with gr.Blocks(css='style.css') as demo:
     if os.getenv('IS_SHARED_UI'):
         show_warning(SHARED_UI_WARNING)
@@ -164,6 +182,10 @@ with gr.Blocks(css='style.css') as demo:
         with gr.TabItem('Test'):
             create_inference_demo(inference_fn)
 
-
-demo.queue(default_enabled=False).launch(share=False)
+    demo.launch(
+        enable_queue=args.enable_queue,
+        server_port=args.port,
+        share=args.share,
+    )
+# demo.queue(default_enabled=False).launch(server_port=args.port, share=args.share)