Update

app.py

@@ -36,7 +36,6 @@ def initialize_game() -> List[List[str]]:
 
     return list(zip(speaker_images, listener_images, targets, roles))
 
-@spaces.GPU(duration=10)
 def get_model_response(
         model, adapter_name, processor, index_to_token, role: str, 
         image_paths: List[str], user_message: str = "", target_image: str = ""
@@ -72,10 +71,10 @@ def get_model_response(
 
 @spaces.GPU(duration=15)
 def get_speaker_response(model, images, input_tokens, attn_mask, image_attn_mask, label, image_paths, processor, img_dir, index_to_token):
+    model = model.cuda()
     with torch.no_grad():
-        print(model.model.device, images.device)
         captions, _, _, _, _ = model.generate(
-            images, input_tokens, attn_mask, image_attn_mask, label,
+            images.cuda(), input_tokens.cuda(), attn_mask.cuda(), image_attn_mask.cuda(), label.cuda(),
             image_paths, processor, img_dir, index_to_token,
             max_steps=30, sampling_type="nucleus", temperature=0.7,
             top_k=50, top_p=1, repetition_penalty=1, num_samples=5
@@ -85,11 +84,12 @@ def get_speaker_response(model, images, input_tokens, attn_mask, image_attn_mask
 @spaces.GPU(duration=15)
 def get_listener_response(model, images, l_input_tokens, l_attn_mask, l_image_attn_mask, index_to_token,
                           s_input_tokens, s_attn_mask, s_image_attn_mask, s_target_mask, s_target_label, image_paths):
+    model = model.cuda()
     with torch.no_grad():
         print(model.model.device, images.device)
         _, _, joint_log_probs = model.comprehension_side([
-            images, l_input_tokens, l_attn_mask, l_image_attn_mask, index_to_token,
-            s_input_tokens, s_attn_mask, s_image_attn_mask, s_target_mask, s_target_label,
+            images.cuda(), l_input_tokens.cuda(), l_attn_mask.cuda(), l_image_attn_mask.cuda(), index_to_token,
+            s_input_tokens.cuda(), s_attn_mask.cuda(), s_image_attn_mask.cuda(), s_target_mask.cuda(), s_target_label.cuda(),
         ])
         target_idx = joint_log_probs[0].argmax().item()
         response = image_paths[target_idx]