Upload 4 files

README.md

@@ -1,10 +1,10 @@
 ---
-title: LLaVA 1.6
+title: LlavaGuard
 emoji: 👁
 colorFrom: green
 colorTo: yellow
 sdk: gradio
-sdk_version: 4.36.1
+sdk_version: 4.16.0
 app_file: app.py
 pinned: false
 license: apache-2.0

app.py

@@ -11,11 +11,11 @@ import gradio as gr
 import spaces
 import torch
 
+from builder import load_pretrained_model
 from llava.constants import IMAGE_TOKEN_INDEX
 from llava.constants import LOGDIR
 from llava.conversation import (default_conversation, conv_templates)
 from llava.mm_utils import KeywordsStoppingCriteria, tokenizer_image_token
-from llava.model.builder import load_pretrained_model
 from llava.utils import (build_logger, violates_moderation, moderation_msg)
 from taxonomy import wrap_taxonomy, default_taxonomy
 
@@ -66,16 +66,16 @@ def run_llava(prompt, pil_image, temperature, top_p, max_new_tokens):
     return outputs[0].strip()
 
 
-# def load_selected_model(model_path):
-#     model_name = model_path.split("/")[-1]
-#     global tokenizer, model, image_processor, context_len
-#     with warnings.catch_warnings(record=True) as w:
-#         warnings.simplefilter("always")
-#         tokenizer, model, image_processor, context_len = load_pretrained_model(model_path, None, model_name)
-#     for warning in w:
-#         if "vision" not in str(warning.message).lower():
-#             print(warning.message)
-#     model.config.tokenizer_model_max_length = 2048 * 2
+def load_selected_model(model_path):
+    model_name = model_path.split("/")[-1]
+    global tokenizer, model, image_processor, context_len
+    with warnings.catch_warnings(record=True) as w:
+        warnings.simplefilter("always")
+        tokenizer, model, image_processor, context_len = load_pretrained_model(model_path, None, model_name)
+    for warning in w:
+        if "vision" not in str(warning.message).lower():
+            print(warning.message)
+    model.config.tokenizer_model_max_length = 2048 * 2
 
 
 def get_conv_log_filename():
@@ -260,11 +260,6 @@ def llava_bot(state, model_selector, temperature, top_p, max_new_tokens, request
         if not os.path.isfile(filename):
             os.makedirs(os.path.dirname(filename), exist_ok=True)
             image.save(filename)
-    # check if model is not None
-    if model is None or tokenizer is None:
-        print(model)
-        print(tokenizer)
-
 
     output = run_llava(prompt, all_images[0], temperature, top_p, max_new_tokens)
 
@@ -507,7 +502,8 @@ Set the environment variable `model` to change the model:
         'LukasHug/LlavaGuard-13B-hf',
         'LukasHug/LlavaGuard-34B-hf', ]
     bits = int(os.getenv("bits", 16))
-    model = os.getenv("model", models[0])
+    model = os.getenv("model", models[1])
+    available_devices = os.getenv("CUDA_VISIBLE_DEVICES", "0")
     model_path, model_name = model, model.split("/")[0]
     if api_key:
         cmd = f"huggingface-cli login --token {api_key} --add-to-git-credential"
@@ -520,7 +516,7 @@ Set the environment variable `model` to change the model:
         model_path = '/common-repos/LlavaGuard/models/LlavaGuard-v1.1-7b-full/smid_and_crawled_v2_with_augmented_policies/json-v16/llava'
 
     print(f"Loading model {model_path}")
-    tokenizer, model, image_processor, context_len = load_pretrained_model(model_path, None, model_name)
+    tokenizer, model, image_processor, context_len = load_pretrained_model(model_path, None, model_name, token=api_key)
     model.config.tokenizer_model_max_length = 2048 * 2
 
     exit_status = 0