Update app.py

app.py

@@ -32,21 +32,9 @@ def print_directory_contents(path, indent=0):
 # 拉取数据集
 os.makedirs(local_dir, exist_ok=True)
 hf_api = HfApi(endpoint="https://huggingface.co", token=hf_token)
-# hf_api = HfApi(token=hf_token)
+hf_api = HfApi(token=hf_token)
 repo_id = "Vchitect/VBench_sampled_video"
 
-downloaded_files = snapshot_download(
-    repo_id=repo_id,
-    token=hf_token,                # 可选：如果你的仓库是私有的
-    repo_type='dataset',           # 如果下载的是数据集而不是模型
-    cache_dir=local_dir,           # 指定缓存/下载位置
-    ignore_patterns=["*.md", "*.txt"]  # 可选：忽略某些模式的文件
-)
-
-print(f"Downloaded files: {downloaded_files}")
-
-model_names = os.listdir(local_dir)
-
 with open("videos_by_dimension.json") as f:
     dimension = json.load(f)['videos_by_dimension']
 
@@ -66,15 +54,31 @@ def get_random_video():
     random_index = random.randint(0, len(model_names) - 1)
     model_name =  model_names[random_index]
 
-    video_path = os.path.join(model_name, type, prompt)
-    if os.path.exists(video_path):
-        print(video_path)
-        return video_path
-    else:
-        video_path = os.path.join(model_name, prompt)
-        if os.path.exists(video_path):
-            print(video_path)
-            return video_path
+    
+
+    video_path_subfolder = os.path.join(model_name, type)
+    try:
+        hf_api.hf_hub_download(
+            repo_id = repo_id,
+            filename = prompt,
+            subfolder = video_path_subfolder,
+            repo_type = dataset,
+            local_dir = local_dir
+        )
+    except Exception as e:
+        print(f"[PATH]{video_path_subfolder} NOT in hf repo, try {model_name}")
+        print(e)
+        video_path_subfolder = model_name
+        try:
+            hf_api.hf_hub_download(
+                repo_id = repo_id,
+                filename = prompt,
+                subfolder = video_path_subfolder,
+                repo_type = dataset,
+                local_dir = local_dir
+            )
+        except Exception as e:
+            print(e)
     # video_path = dataset['train'][random_index]['video_path']
     print('error:', video_path)
     return video_path