Update

app.py

@@ -42,40 +42,59 @@ def get_model_response(
         image_paths: List[str], user_message: str = "", target_image: str = ""
 ) -> str:    
     model.model.set_adapter(adapter_name)
-    print(model.model.active_adapter)
     if role == "speaker":
         img_dir = "tangram_pngs"
+        print("Starting processing")
         input_tokens, attn_mask, images, image_attn_mask, label = joint_speaker_input(
             processor, image_paths, target_image, model.get_listener().device
         )
-        print("Hi")
-        with torch.no_grad():
-            image_paths = [image_paths]
-            captions, _, _, _, _ = model.generate(
-                images, input_tokens, attn_mask, image_attn_mask, label,
-                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
-            )
-        print("There")
+        image_paths = [image_paths]
+        print("Starting inference")
+        captions = get_speaker_response(model, images, input_tokens, attn_mask, image_attn_mask, label, image_paths,
+                                        processor, img_dir, index_to_token)
+        print("Done")
         response = captions[0]
     else:  # listener
+        print("Starting processing")
         images, l_input_tokens, l_attn_mask, l_image_attn_mask, s_input_tokens, s_attn_mask, \
             s_image_attn_mask, s_target_mask, s_target_label = joint_listener_input(
                 processor, image_paths, user_message, model.get_listener().device
             )
 
-        with torch.no_grad():
-            # Forward
-            _, _, 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,
-            ])
-            target_idx = joint_log_probs[0].argmax().item()
-            response = image_paths[target_idx]
+        print("Starting inference")
+        response = 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
+        )
+        print("Done")
 
     return 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):
+    with torch.no_grad():
+        print(model.model.device, images.device)
+        captions, _, _, _, _ = model.generate(
+            images, input_tokens, attn_mask, image_attn_mask, label,
+            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
+        )
+    return captions
+
+@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):
+    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,
+        ])
+        target_idx = joint_log_probs[0].argmax().item()
+        response = image_paths[target_idx]
+    return response            
+
 def interaction(model, processor, index_to_token, model_iteration: str) -> Tuple[List[str], List[str]]:
     image_role_pairs = initialize_game()
     conversation = []