Update app.py

app.py

@@ -1,5 +1,4 @@
 import json
-import os
 import os.path as osp
 import threading
 
@@ -12,10 +11,11 @@ from llava.model.builder import load_pretrained_model
 from llava_utils import prompt_wrapper, generator
 from utils import normalize, denormalize, load_image
 
-
+BASE_DIR = "/workingdir/soh337/llavaguard"
+# BASE_DIR = "/Users/ahren/Workspace/Multimodal/llavaguard"
 UNCONSTRAINED_ATTACK_IMAGE_PATH = 'unconstrained_attack_images/adversarial_'
 CONSTRAINED_ATTACK_IMAGE_PATH = 'adversarial_qna_images/adv_image_'
-MODEL_PATH = "liuhaotian/llava-v1.5-13b"
+MODEL_PATH = "/workingdir/models_hf/liuhaotian/llava-v1.5-13b"
 
 TEXT_SAFETY_PATCHES = {
     "optimized": "text_patch_optimized",
@@ -50,11 +50,16 @@ def load_model_async(model_path, model_name):
 
     my_generator = generator.Generator(model=model, tokenizer=tokenizer)
 
-threading.Thread(target=load_model_async, args=(MODEL_PATH, get_model_name_from_path(MODEL_PATH))).start()
+# threading.Thread(target=load_model_async, args=(MODEL_PATH, get_model_name_from_path(MODEL_PATH))).start()
+
+
 
 print('>>> Initializing Models')
 
-prompts = rtp_read(osp.join('harmful_corpus/rtp_prompts.jsonl'))
+load_model_async(MODEL_PATH, get_model_name_from_path(MODEL_PATH))
+
+
+prompts = rtp_read(osp.join(BASE_DIR, 'harmful_corpus/rtp_prompts.jsonl'))
 
 # out_unprotected: responses without the safety patch
 out, out_unprotected = [], []
@@ -248,6 +253,5 @@ with gr.Blocks(css=css) as demo:
                 gr.Textbox(label="NO Safety Patches")
             ])
 
-
 # Launch the demo
-demo.launch()
+demo.launch()