hftk

app.py

@@ -57,7 +57,7 @@ if __name__ == "__main__":
                 task_btn = gr.Button("See current task status")
                 once_btn = gr.Button("Trigger once manually")
                 smtp_btn = gr.Button("SMTP test")
-                with gr.Tab("HF cookie test"):
+                with gr.Tab("HF token test"):
                     hf_txt = gr.Textbox(
                         label="Restart a space with permissions",
                         placeholder="username/repo",
@@ -65,7 +65,7 @@ if __name__ == "__main__":
                     )
                     with gr.Row():
                         clhf_btn = gr.Button("Clear")
-                        hf_btn = gr.Button("Restart to test cookie validity")
+                        hf_btn = gr.Button("Restart to test token validity")
 
                 with gr.Tab("MS cookie test"):
                     ms_txt = gr.Textbox(

modules/restart.py

@@ -1,15 +1,18 @@
 import requests
+from huggingface_hub import HfApi
 from modules.smtp import send_email
 from utils import (
     HF_DOMAIN,
     MS_DOMAIN,
-    HF_HEADER,
+    HF_TK,
     MS_HEADER,
     TIMEOUT,
     PUB_REPOS,
     PRI_REPOS,
 )
 
+HF_API = HfApi(token=HF_TK)
+
 
 def restart_private_studio(repo: str):
     response = requests.put(
@@ -21,16 +24,6 @@ def restart_private_studio(repo: str):
     return f"\n{response.text}\n"
 
 
-def restart_private_space(repo: str):
-    response = requests.post(
-        f"{HF_DOMAIN}/api/spaces/{repo}/restart",
-        headers=HF_HEADER,
-        timeout=TIMEOUT,
-    )
-    response.raise_for_status()
-    return f"\n{response.text}\n"
-
-
 def restart_repos():
     if PUB_REPOS:
         repos = PUB_REPOS.split(";")
@@ -38,7 +31,7 @@ def restart_repos():
             private_repo = repo.strip()
             print(f"Restarting {private_repo} space & studio...")
             restart_private_studio(private_repo)
-            restart_private_space(private_repo)
+            HF_API.restart_space(private_repo)
 
 
 def self_restart():
@@ -46,7 +39,7 @@ def self_restart():
         if PRI_REPOS:
             repos = PRI_REPOS.split(";")
             for repo in repos:
-                restart_private_space(repo)
+                HF_API.restart_space(repo)
 
     except Exception as e:
         send_email(f"Failed to self-restart: {e}")
@@ -64,7 +57,7 @@ def test_hf_restart(repo: str):
             repo = repo.split("/spaces/")[1].split("/")[:2]
             repo = "/".join(repo)
 
-        logs = restart_private_space(repo)
+        logs = HF_API.restart_space(repo)
         logs += f"\nSuccessfully restart {HF_DOMAIN}/spaces/{repo}\n"
 
     except Exception as e:

utils.py

@@ -7,7 +7,7 @@ PRI_REPOS = os.getenv("pri_repos")
 API = os.getenv("api")
 EMAIL = os.getenv("email")
 SMTP = os.getenv("smtp")
-HF_CK = os.getenv("hfck")
+HF_TK = os.getenv("hftk")
 MS_CK = os.getenv("msck")
 TAG = os.getenv("target")
 if not TAG:
@@ -29,10 +29,6 @@ MS_DOMAIN = "https://www.modelscope.cn"
 HEADER = {
     "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537"
 }
-HF_HEADER = {
-    "User-Agent": HEADER["User-Agent"],
-    "Cookie": HF_CK,
-}
 MS_HEADER = {
     "User-Agent": HEADER["User-Agent"],
     "Cookie": MS_CK,