Update app.py

app.py

@@ -33,6 +33,13 @@
 # --------------------------------------------------------
 # gradio demo executable
 # --------------------------------------------------------
+#!/usr/bin/env python3
+# Copyright (C) 2024-present Naver Corporation. All rights reserved.
+# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
+#
+# --------------------------------------------------------
+# gradio demo executable
+# --------------------------------------------------------
 import os
 import torch
 import tempfile
@@ -46,17 +53,24 @@ from mast3r.utils.misc import hash_md5
 import matplotlib.pyplot as pl
 pl.ion()
 
-torch.backends.cuda.matmul.allow_tf32 = True  # for gpu >= Ampere and pytorch >= 1.12
+torch.backends.cuda.matmul.allow_tf32 = True  # for GPU >= Ampere and PyTorch >= 1.12
+
+def get_default_weights_path(model_name):
+    # Construct default weights path based on model_name
+    return f"naver/{model_name}"
 
 if __name__ == '__main__':
     parser = get_args_parser()
     args = parser.parse_args()
 
-    # Set default value for `args.weights` if not provided
+    # Ensure at least one of weights or model_name is provided
+    if args.weights is None and args.model_name is None:
+        # Provide a default model_name if both weights and model_name are not provided
+        args.model_name = 'MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric'
+    
+    # If weights are not provided but model_name is, construct weights_path
     if args.weights is None:
-        # Set a default model_name if weights are not provided
-        args.model_name = 'MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric'  # Default model_name
-        args.weights = "naver/" + args.model_name  # Construct default weights path
+        args.weights = get_default_weights_path(args.model_name)
 
     if args.server_name is not None:
         server_name = args.server_name
@@ -66,12 +80,14 @@ if __name__ == '__main__':
     # Use the provided or default weights_path
     weights_path = args.weights
 
+    # Load the model with the weights_path
     model = AsymmetricMASt3R.from_pretrained(weights_path).to(args.device)
     chkpt_tag = hash_md5(weights_path)
 
     def get_context(tmp_dir):
         return tempfile.TemporaryDirectory(suffix='_mast3r_gradio_demo') if tmp_dir is None \
             else nullcontext(tmp_dir)
+    
     with get_context(args.tmp_dir) as tmpdirname:
         cache_path = os.path.join(tmpdirname, chkpt_tag)
         os.makedirs(cache_path, exist_ok=True)