Update app.py

app.py

@@ -3,14 +3,11 @@ import os
 import argparse
 import time
 import subprocess
-if '/workspace' not in sys.path:
-    sys.path.append('/workspace')
-import gradio_web_server as gws
-# from llavaguard.hf_utils import set_up_env_and_token
 
+import gradio_web_server as gws
 
 # Execute the pip install command with additional options
-# subprocess.check_call([sys.executable, '-m', 'pip', 'install', 'flash-attn', '--no-build-isolation', '-U'])
+subprocess.check_call([sys.executable, '-m', 'pip', 'install', 'flash-attn', '--no-build-isolation', '-U'])
 
 
 def start_controller():
@@ -84,6 +81,15 @@ Set the environment variable `model` to change the model:
     # set the huggingface login token
     controller_proc = start_controller()
     concurrency_count = int(os.getenv("concurrency_count", 5))
+    api_key = os.getenv("token")
+    if api_key:
+        cmd = f"huggingface-cli login --token {api_key} --add-to-git-credential"
+        os.system(cmd)
+    else:
+        if '/workspace' not in sys.path:
+            sys.path.append('/workspace')
+        from llavaguard.hf_utils import set_up_env_and_token
+        set_up_env_and_token(read=True, write=False)
 
     models = [
         'LukasHug/LlavaGuard-7B-hf',