Update app.py

app.py

@@ -115,21 +115,33 @@ def predict(image_input, question):
         encoded = text_tokenizer(prompt, return_tensors="pt").to(device)
 
         with torch.no_grad():
-            logits = model(image, encoded["input_ids"], encoded["attention_mask"])
-
-        # Generate answer
-        generated_tokens = model.phi3.generate(
-            inputs=None, # Remove input_ids and attention_mask
-            inputs_embeds=logits, # Use the logits from the forward pass as input embeddings
-            max_length=128,
-            pad_token_id=text_tokenizer.eos_token_id
-        )
+            # Get image embeddings
+            image_embeddings = model.image_encoder(image)
+            projected_image_embeddings = model.image_projection(image_embeddings)
+
+            # Reshape image embeddings to (batch_size, 1, phi3_embed_dim)
+            projected_image_embeddings = projected_image_embeddings.unsqueeze(1)
+
+            # Concatenate along the sequence dimension (dim=1)
+            extended_attention_mask = torch.cat([torch.ones(projected_image_embeddings.shape[:2], device=encoded["attention_mask"].device), encoded["attention_mask"]], dim=1)
+            extended_input_ids = torch.cat([torch.zeros(projected_image_embeddings.shape[:2], dtype=torch.long, device=encoded["input_ids"].device), encoded["input_ids"]], dim=1)
+
+            # Generate answer
+            generated_tokens = model.phi3.generate(
+                input_ids=extended_input_ids,
+                attention_mask=extended_attention_mask,
+                max_length=200,
+                pad_token_id=text_tokenizer.eos_token_id
+            )
 
         answer = text_tokenizer.decode(generated_tokens[0], skip_special_tokens=True)
         answer = answer.replace(prompt, "").strip() # Remove prompt from answer
 
         return answer
 
+    except Exception as e:
+        return f"An error occurred: {str(e)}"
+
     except Exception as e:
         #return f"An error occurred: {str(e)}"
         return f"An error occurred: {traceback.format_exc()}"