Add story display

app.py

@@ -26,11 +26,15 @@ if uploaded_file is not None:
     enc_image = model.encode_image(image)
     answer = model.answer_question(enc_image, "Describe this image in a few sentences.", tokenizer)
     
-    st.write(answer)
+    st.header("Image Description Generated.")
+    st.success(answer)
 
     model_id = "meta-llama/Meta-Llama-3-8B"
 
     pipeline = transformers.pipeline("text-generation", model=model_id, model_kwargs={"torch_dtype": torch.bfloat16}, device_map="auto")
 
-    pipeline("Generate a short story based on image description: " + answer)
+    story = pipeline("Generate a short story based on image description: " + answer)
+
+    st.header("Story Generated.")
+    st.success(story)