diff --git a/.github/workflows/docker-build.yml b/.github/workflows/docker-build.yml
index 180af936e..f480d3214 100644
--- a/.github/workflows/docker-build.yml
+++ b/.github/workflows/docker-build.yml
@@ -100,6 +100,4 @@ jobs:
 
       - name: Restart HuggingFace Spaces Build
         run: |
-          poetry run python ./scripts/factory_restart_space.py
-        env:
-          HUGGINGFACE_API_TOKEN: ${{ secrets.HUGGINGFACE_API_TOKEN }}
+          poetry run python ./scripts/factory_restart_space.py --space "Langflow/Langflow-Preview" --token ${{ secrets.HUGGINGFACE_API_TOKEN }}
diff --git a/scripts/factory_restart_space.py b/scripts/factory_restart_space.py
index e9972e8cb..8a52dc42a 100644
--- a/scripts/factory_restart_space.py
+++ b/scripts/factory_restart_space.py
@@ -1,4 +1,4 @@
-import os
+import argparse
 
 from huggingface_hub import HfApi, list_models
 from rich import print
@@ -6,11 +6,27 @@ from rich import print
 # Use root method
 models = list_models()
 
+args = argparse.ArgumentParser(description="Restart a space in the Hugging Face Hub.")
+args.add_argument("--space", type=str, help="The space to restart.")
+args.add_argument("--token", type=str, help="The Hugging Face API token.")
+
+parsed_args = args.parse_args()
+
+space = parsed_args.space
+
+if not space:
+    print("Please provide a space to restart.")
+    exit()
+
+if not parsed_args.api_token:
+    print("Please provide an API token.")
+    exit()
+
 # Or configure a HfApi client
 hf_api = HfApi(
     endpoint="https://huggingface.co",  # Can be a Private Hub endpoint.
-    token=os.getenv("HUGGINFACE_API_TOKEN"),
+    token=parsed_args.token,
 )
 
-space_runtime = hf_api.restart_space("Langflow/Langflow-Preview", factory_reboot=True)
+space_runtime = hf_api.restart_space(space, factory_reboot=True)
 print(space_runtime)