Update app.py

app.py

@@ -12,7 +12,6 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 # Fixed prompt
 FIXED_PROMPT = "Generate a funny caption from image"
 
-# Function to generate caption for the uploaded image with the fixed prompt
 # Function to generate caption for the uploaded image with the fixed prompt
 def generate_caption(image):
     # Preprocess the image
@@ -33,7 +32,11 @@ def generate_caption(image):
     batch_size = image_tensor.shape[0]
     repeated_prompt_tensor = fixed_prompt_tensor.repeat(batch_size, 1)
     
-    # Concatenate the prompt tensor with the image tensor along the sequence dimension
+    # Reshape the image tensor to match the shape of the prompt tensor
+    # The reshaping depends on the model's input requirements
+    image_tensor = image_tensor.view(batch_size, -1)
+    
+    # Concatenate the prompt tensor with the reshaped image tensor along the sequence dimension
     input_tensor = torch.cat((repeated_prompt_tensor, image_tensor), dim=1)
     
     # Generate caption
@@ -42,6 +45,7 @@ def generate_caption(image):
     return caption
 
 
+
 # Streamlit app
 st.title("Cartoon Caption Generator")