Update app.py

app.py

@@ -28,15 +28,24 @@ if uploaded_image is not None:
 
     # Generate story from caption
     if st.button('Generate Story from Caption'):
-        st.write("Generating story...")
-        story_prompt = f"Based on the image, here's a story: {caption}"
-
-        # Encode prompt text to input ids
-        input_ids = tokenizer_gpt2.encode(story_prompt, return_tensors='pt')
-
-        # Generate text using GPT-2 model
-        story = model_gpt2.generate(input_ids, max_length=200, num_return_sequences=1, pad_token_id=tokenizer_gpt2.eos_token_id)
-
-        # Decode and display the story
-        story_text = tokenizer_gpt2.decode(story[0], skip_special_tokens=True)
-        st.text_area("Generated Story", story_text, height=250)
+      st.write("Generating story...")
+      story_prompt = f"Based on the image, here's a story: {caption}"
+
+    # Encode prompt text to input ids
+    input_ids = tokenizer_gpt2.encode(story_prompt, return_tensors='pt')
+
+    # Generate text using GPT-2 model
+    story = model_gpt2.generate(
+        input_ids,
+        max_length=200,
+        num_return_sequences=1,
+        no_repeat_ngram_size=2,
+        temperature=0.9,
+        top_k=50,
+        top_p=0.95,
+        pad_token_id=tokenizer_gpt2.eos_token_id
+    )
+
+    # Decode and display the story
+    story_text = tokenizer_gpt2.decode(story[0], skip_special_tokens=True)
+    st.text_area("Generated Story", story_text, height=250)