Update app.py

app.py

@@ -22,7 +22,14 @@ def generate_story(image, theme, genre):
         image_features = feature_extractor(images=image, return_tensors="pt")
         
         # Generate image caption
-        caption_ids = model.generate(image_features.pixel_values, max_length=50, num_beams=3, temperature=1.0)
+        # Generate image caption
+        caption_ids = model.generate(
+            input_ids=image_features.pixel_values,
+            max_length=50,
+            num_beams=3,
+            temperature=1.0,
+            attention_mask=image_features.attention_mask  # Pass the attention_mask
+        )
         
         # Decode the caption
         caption_text = tokenizer.batch_decode(caption_ids, skip_special_tokens=True)[0]