Upload 数据库连接.py

数据库连接.py

@@ -6,12 +6,21 @@ from huggingface_hub import HfApi, hf_hub_download
 HF_TOKEN = os.environ.get("HF_TOKEN")
 # 已经替换为你真实的 Dataset 仓库
 DATASET_REPO_ID = "ZHIWEI666/ComfyUI-Ranking" 
-# HF Spaces 的 Docker 环境只有 /tmp 目录允许自由读写
-LOCAL_DB_DIR = "/tmp/local_db_data" 
+
+# ==========================================
+# 【核心修改】：智能判断运行环境
+# ==========================================
+if os.environ.get("SPACE_ID"):
+    # 如果检测到在 Hugging Face Spaces 环境运行，则使用 /tmp 目录（绕过容器只读限制）
+    LOCAL_DB_DIR = "/tmp/local_db_data"
+else:
+    # 本地运行时，自动在当前 Python 文件同级目录下创建一个 "cache" 文件夹
+    BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+    LOCAL_DB_DIR = os.path.join(BASE_DIR, "cache")
 
 api = HfApi() if HF_TOKEN else None
 
-# 确保本地测试目录存在
+# 确保本地缓存目录存在
 if not os.path.exists(LOCAL_DB_DIR):
     os.makedirs(LOCAL_DB_DIR)
 
@@ -38,7 +47,7 @@ def save_data(file_name: str, data):
     """保存数据：写入本地并同步到 HF Dataset"""
     local_path = os.path.join(LOCAL_DB_DIR, file_name)
     
-    # 1. 始终先写入本地临时目录
+    # 1. 始终先写入本地缓存目录
     with open(local_path, "w", encoding="utf-8") as f:
         json.dump(data, f, ensure_ascii=False, indent=2)
     
@@ -56,7 +65,7 @@ def save_data(file_name: str, data):
         except Exception as e:
             print(f"同步到 HF Dataset 失败: {e}")
 
-# --- 【新增】：保存真实的二进制文件 ---
+# --- 保存真实的二进制文件 ---
 def save_file(file_path_in_repo: str, content: bytes):
     local_full_path = os.path.join(LOCAL_DB_DIR, file_path_in_repo)