Refactor model loading to use snapshot_download from Hugging Face Hub and streamline file management for pretrained models.

app.py

@@ -7,7 +7,7 @@ import argparse
 import gradio as gr
 import uuid
 import spaces
-from huggingface_hub import hf_hub_download
+from huggingface_hub import snapshot_download
 #
 
 subprocess.run(shlex.split("pip install wheel/torch_scatter-2.1.2+pt21cu121-cp310-cp310-linux_x86_64.whl"))
@@ -20,50 +20,27 @@ subprocess.run(shlex.split("pip install wheel/pointops-1.0-cp310-cp310-linux_x86
 from src.utils.visualization_utils import render_video_from_file
 from src.model import LSM_MASt3R
 
-# 定义相对路径和 Hugging Face 仓库信息
-relative_model_dir = "checkpoints"  # 文件夹名称
-relative_model_path = os.path.join(relative_model_dir, "checkpoint-40.pth")  # 相对路径
-model_repo = "kairunwen/LSM"  # Hugging Face 仓库
-model_filename = "checkpoint-40.pth"  # 仓库中的文件名
-
-# 转换为绝对路径
-model_path = os.path.abspath(relative_model_path)
-
-try:
-    # 创建 checkpoints 文件夹（如果不存在）
-    os.makedirs(relative_model_dir, exist_ok=True)
-    print(f"确保 {relative_model_dir} 文件夹存在")
-
-    # 验证文件是否存在
-    if os.path.exists(model_path):
-        print(f"找到本地模型文件: {model_path}")
-    else:
-        print(f"本地模型文件 {model_path} 不存在，正在从 Hugging Face 下载...")
-        model_path = hf_hub_download(repo_id=model_repo, filename=model_filename)
-        # 可选：将下载的文件移动到 checkpoints 文件夹
-        import shutil
-        shutil.move(model_path, os.path.abspath(relative_model_path))
-        model_path = os.path.abspath(relative_model_path)
-        print(f"模型文件已下载并移动到: {model_path}")
-
-    # 加载模型
-    model = LSM_MASt3R.from_pretrained(model_path, device='cuda')
-    model = model.eval()
-    print("模型加载成功并设置为评估模式！")
-
-except FileNotFoundError:
-    print(f"错误: 无法找到或下载文件 {model_filename}，请检查路径或仓库 {model_repo}。")
-except KeyError as e:
-    print(f"错误: 检查点文件格式不正确，缺少键 {e}。请确认 checkpoint-40.pth 包含 'args' 和 'model'。")
-except Exception as e:
-    print(f"发生未知错误: {e}")
-    # 调试：检查检查点内容
-    ckpt = torch.load(model_path, map_location='cpu')
-    print("检查点键:", ckpt.keys())
-    print("config.model:", ckpt['args'].model)
-
-
-
+# Download the model checkpoint from Hugging Face Hub
+repo_id = "Journey9ni/LSM"
+remote_dir = "checkpoints/pretrained_models/"
+local_dir = "checkpoints/pretrained_model"
+model_path_map = {
+    "MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric.pth": "MASt3R_ViTLarge_BaseDecoder_512_catmlpdpt_metric.pth",
+    "checkpoint-40.pth":"checkpoint-40.pth",
+    "demo_e200.ckpt":"lang_seg.ckpt"
+}
+os.makedirs(local_dir, exist_ok=True)
+# download remote repo
+snapshot_download(repo_id=repo_id)
+
+# rename the files
+for remote_name, local_name in model_path_map.items():
+    os.rename(os.path.join(local_dir, remote_name), os.path.join(local_dir, local_name))
+
+# load the model
+model_path = "checkpoints/pretrained_model/checkpoint-40.pth"
+model = LSM_MASt3R.from_pretrained(model_path, device='cuda')
+model = model.eval()
 
 @spaces.GPU(duration=80)
 def process(inputfiles, input_path=None):