space_waker / app.py
John6666's picture
Upload app.py
04475d4 verified
import gradio as gr
from huggingface_hub import HfApi
import os
HF_TOKEN = os.environ.get("HF_TOKEN")
HF_USER = os.environ.get("HF_USER")
def wakeup_spaces(username: str, hf_token: str="", is_private=True, progress=gr.Progress(track_tqdm=True)):
try:
api = HfApi(token=hf_token if hf_token else HF_TOKEN)
space_infos = api.list_spaces(author=username)
spaces = []
logs = []
for s in space_infos:
if s.private and not is_private: continue
spaces.append(s.id)
for s in spaces:
stage = api.get_space_runtime(repo_id=s).stage
if stage in ("RUNNING", "APP_STARTING", "BUILDING"):
print(f"'{s}' is running.")
progress(0, desc=f"'{s}' is running.")
else:
print(f"'{s}' isn't running. Restarting...")
progress(0, desc=f"'{s}' isn't running. Restarting...")
api.restart_space(repo_id=s)
logs.append(f"'{s}' is in '{stage}'.")
logs.append(f"'{s}' is restarted.")
return "<br>\n".join(logs)
except Exception as e:
print(f"Error: {e}")
raise gr.Error(f"Error: {e}")
with gr.Blocks() as demo:
with gr.Row():
username = gr.Textbox(label="HF Username", value=HF_USER, lines=1)
token = gr.Textbox(label="HF Write Token", value="", lines=1)
is_private = gr.Checkbox(label="Apply to private space", value=True)
run_button = gr.Button("Wake up spaces", variant="primary")
info_md = gr.Markdown("<br><br><br>")
run_button.click(wakeup_spaces, [username, token, is_private], [info_md])
demo.launch()