Update demo.

app.py

@@ -61,24 +61,23 @@ The service is a research preview intended for non-commercial use only, subject
 
 
 class Chat:
-    def __init__(self, model_path, conv_mode, model_base=None, load_8bit=False, load_4bit=False, device='cuda'):
+    def __init__(self, model_path, conv_mode, model_base=None, load_8bit=False, load_4bit=False):
         # disable_torch_init()
         model_name = get_model_name_from_path(model_path)
         self.tokenizer, self.model, processor, context_len = load_pretrained_model(
             model_path, model_base, model_name,
             load_8bit, load_4bit,
-            device=device,
             offload_folder="save_folder")
         self.processor = processor
         self.conv_mode = conv_mode
         self.conv = conv_templates[conv_mode].copy()
-        self.device = self.model.device
 
     def get_prompt(self, qs, state):
         state.append_message(state.roles[0], qs)
         state.append_message(state.roles[1], None)
         return state
 
+    @spaces.GPU(duration=120)
     @torch.inference_mode()
     def generate(self, tensor: list, modals: list, prompt: str, first_run: bool, state):
         # TODO: support multiple turns of conversation.
@@ -92,7 +91,7 @@ class Chat:
         prompt = state.get_prompt()
         # print('\n\n\n')
         # print(prompt)
-        input_ids = tokenizer_MMODAL_token(prompt, tokenizer, MMODAL_TOKEN_INDEX[modals[0]], return_tensors='pt').unsqueeze(0).to(self.device)
+        input_ids = tokenizer_MMODAL_token(prompt, tokenizer, MMODAL_TOKEN_INDEX[modals[0]], return_tensors='pt').unsqueeze(0).to(self.model.device)
 
         # 3. generate response according to visual signals and prompts. 
         stop_str = self.conv.sep if self.conv.sep_style in [SeparatorStyle.SINGLE] else self.conv.sep2