Update web_demo.py

web_demo.py

@@ -4,12 +4,14 @@ import numpy as np
 import copy
 import gradio as gr
 import sys
+
+import spaces
+
 from vita_audio.tokenizer import get_audio_tokenizer
 from vita_audio.data.processor.audio_processor import add_audio_input_contiguous
 
 
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, AutoConfig
-from transformers.generation import GenerationConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, AutoConfig, GenerationConfig
 
 
 
@@ -67,6 +69,8 @@ def _parse_text(text):
 
 
 def _launch_demo(model, tokenizer, audio_tokenizer):
+
+    @spaces.GPU(duration=120)
     def predict(_chatbot, task_history,task):
         chat_query = task_history[-1][0]
         print(task_history)