LukasHug commited on
Commit
1342a94
1 Parent(s): 62b8fab

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -3
app.py CHANGED
@@ -11,11 +11,11 @@ import gradio as gr
11
  import spaces
12
  import torch
13
 
14
- from builder import load_pretrained_model
15
  from llava.constants import IMAGE_TOKEN_INDEX
16
  from llava.constants import LOGDIR
17
  from llava.conversation import (default_conversation, conv_templates)
18
  from llava.mm_utils import KeywordsStoppingCriteria, tokenizer_image_token
 
19
  from llava.utils import (build_logger, violates_moderation, moderation_msg)
20
  from taxonomy import wrap_taxonomy, default_taxonomy
21
 
@@ -502,7 +502,7 @@ Set the environment variable `model` to change the model:
502
  'LukasHug/LlavaGuard-13B-hf',
503
  'LukasHug/LlavaGuard-34B-hf', ]
504
  bits = int(os.getenv("bits", 16))
505
- model = os.getenv("model", models[0])
506
  available_devices = os.getenv("CUDA_VISIBLE_DEVICES", "0")
507
  model_path, model_name = model, model.split("/")[0]
508
  if api_key:
@@ -516,7 +516,7 @@ Set the environment variable `model` to change the model:
516
  model_path = '/common-repos/LlavaGuard/models/LlavaGuard-v1.1-7b-full/smid_and_crawled_v2_with_augmented_policies/json-v16/llava'
517
 
518
  print(f"Loading model {model_path}")
519
- tokenizer, model, image_processor, context_len = load_pretrained_model(model_path, None, model_name, token=api_key)
520
  model.config.tokenizer_model_max_length = 2048 * 2
521
 
522
  exit_status = 0
 
11
  import spaces
12
  import torch
13
 
 
14
  from llava.constants import IMAGE_TOKEN_INDEX
15
  from llava.constants import LOGDIR
16
  from llava.conversation import (default_conversation, conv_templates)
17
  from llava.mm_utils import KeywordsStoppingCriteria, tokenizer_image_token
18
+ from llava.model.builder import load_pretrained_model
19
  from llava.utils import (build_logger, violates_moderation, moderation_msg)
20
  from taxonomy import wrap_taxonomy, default_taxonomy
21
 
 
502
  'LukasHug/LlavaGuard-13B-hf',
503
  'LukasHug/LlavaGuard-34B-hf', ]
504
  bits = int(os.getenv("bits", 16))
505
+ model = os.getenv("model", models[1])
506
  available_devices = os.getenv("CUDA_VISIBLE_DEVICES", "0")
507
  model_path, model_name = model, model.split("/")[0]
508
  if api_key:
 
516
  model_path = '/common-repos/LlavaGuard/models/LlavaGuard-v1.1-7b-full/smid_and_crawled_v2_with_augmented_policies/json-v16/llava'
517
 
518
  print(f"Loading model {model_path}")
519
+ tokenizer, model, image_processor, context_len = load_pretrained_model(model_path, None, model_name)
520
  model.config.tokenizer_model_max_length = 2048 * 2
521
 
522
  exit_status = 0