jedyang97 commited on
Commit
252ad67
1 Parent(s): 30728d3
Files changed (2) hide show
  1. app.py +1 -0
  2. model.py +3 -1
app.py CHANGED
@@ -46,6 +46,7 @@ tokenizer, model, data_loader = load_model_and_dataloader(
46
  scene_to_obj_mapping=scene_to_obj_mapping,
47
  )
48
 
 
49
  def get_chatbot_response(user_chat_input, scene_id):
50
  # Get the response from the model
51
  prompt, response = get_model_response(
 
46
  scene_to_obj_mapping=scene_to_obj_mapping,
47
  )
48
 
49
+ @spaces.GPU
50
  def get_chatbot_response(user_chat_input, scene_id):
51
  # Get the response from the model
52
  prompt, response = get_model_response(
model.py CHANGED
@@ -1,4 +1,5 @@
1
  # model.py
 
2
  import json
3
  import torch
4
  from torch.utils.data import DataLoader
@@ -11,7 +12,7 @@ from llava.mm_utils import get_model_name_from_path
11
  from llava.model.builder import load_pretrained_model
12
  from llava.utils import disable_torch_init
13
 
14
-
15
  def load_model_and_dataloader(model_path, model_base, scene_to_obj_mapping, obj_context_feature_type="text", load_8bit=False, load_4bit=False, load_bf16=False):
16
  disable_torch_init()
17
 
@@ -41,6 +42,7 @@ def load_model_and_dataloader(model_path, model_base, scene_to_obj_mapping, obj_
41
  return tokenizer, model, data_loader
42
 
43
 
 
44
  def get_model_response(model, tokenizer, data_loader, scene_id, user_input, max_new_tokens=50, temperature=0.2, top_p=0.9):
45
  input_data = [
46
  {
 
1
  # model.py
2
+ import spaces
3
  import json
4
  import torch
5
  from torch.utils.data import DataLoader
 
12
  from llava.model.builder import load_pretrained_model
13
  from llava.utils import disable_torch_init
14
 
15
+ @spaces.GPU
16
  def load_model_and_dataloader(model_path, model_base, scene_to_obj_mapping, obj_context_feature_type="text", load_8bit=False, load_4bit=False, load_bf16=False):
17
  disable_torch_init()
18
 
 
42
  return tokenizer, model, data_loader
43
 
44
 
45
+ @spaces.GPU
46
  def get_model_response(model, tokenizer, data_loader, scene_id, user_input, max_new_tokens=50, temperature=0.2, top_p=0.9):
47
  input_data = [
48
  {