Use model to generate story based on image description

app.py

@@ -1,3 +1,5 @@
+import transformers
+import torch
 import einops
 import streamlit as st
 from PIL import Image
@@ -25,3 +27,10 @@ if uploaded_file is not None:
     answer = model.answer_question(enc_image, "Describe this image in a few sentences.", tokenizer)
     
     st.write(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)
+