Update app.py

app.py

@@ -35,6 +35,19 @@ TOKEN = "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"
 
 pipe = pipeline("translation", model="Helsinki-NLP/opus-mt-en-es")
 
+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('--live', action='store_true')
+    parser.add_argument('--share', action='store_true')
+    parser.add_argument('--port', type=int)
+    parser.add_argument('--disable-queue',
+                        dest='enable_queue',
+                        action='store_false')
+    parser.add_argument('--allow-flagging', type=str, default='never')
+    return parser.parse_args()
+
 def load_model(file_name: str, path:str,device: torch.device) -> nn.Module:
     path = hf_hub_download(f'{path}',
                            f'{file_name}',
@@ -59,7 +72,10 @@ def main():
     #else:
     #    ini = "False1"
     #result = subprocess.check_output(['nvidia-smi'])
-    load_model("stylegan2-ffhq-config-f","feng2022/Time-TravelRephotography_stylegan2-ffhq-config-f")
+    
+    args = parse_args()
+    device = torch.device(args.device)
+    load_model("stylegan2-ffhq-config-f","feng2022/Time-TravelRephotography_stylegan2-ffhq-config-f",device)
     
     iface = gr.Interface(
       fn=predict,