Update tasks/eval/model_utils.py

tasks/eval/model_utils.py

@@ -8,6 +8,7 @@ from tasks.eval.eval_utils import Conversation
 from models.pllava import PllavaProcessor, PllavaForConditionalGeneration, PllavaConfig
 from accelerate import init_empty_weights, dispatch_model, infer_auto_device_map,load_checkpoint_in_model
 from accelerate.utils import get_balanced_memory
+import spaces
 
 from transformers import StoppingCriteria
 class KeywordsStoppingCriteria(StoppingCriteria):
@@ -136,6 +137,7 @@ def load_adapters(model, adapter_model_name_or_paths):
     return model
 
 
+@spaces.GPU(duration=120)
 def pllava_answer(conv: Conversation, model, processor, img_list, do_sample=True, max_new_tokens=200, num_beams=1, min_length=1, top_p=0.9,
                repetition_penalty=1.0, length_penalty=1, temperature=1.0, stop_criteria_keywords=None, print_res=False):
     # torch.cuda.empty_cache()