Update app.py

app.py

@@ -2,6 +2,8 @@ import streamlit as st
 from transformers import pipeline
 import torch
 from transformers import pipeline
+from diffusers import FluxPipeline
+
 
 model_id = "meta-llama/Llama-3.2-3B-Instruct"
 pipe = pipeline(
@@ -22,8 +24,24 @@ def poet(text):
     )
     print(outputs[0]["generated_text"][-1])
 
+def poet_image(poetry):
+    flux = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
+    flux.enable_model_cpu_offload()
+    image = flux(
+        poetry,
+        height=1024,
+        width=1024,
+        guidance_scale=3.5,
+        num_inference_steps=50,
+        max_sequence_length=512,
+        generator=torch.Generator("cpu").manual_seed(0)
+    ).images[0]
+    return image
+
 st.title("Shakespeare Ai")
 st.write("A space made to allow people to create shakespeare like text with images!")
 prompt = st.text_input("Enter your prompt: ")
 st.button("Generate Shakespeare")
-st.write(poet(prompt))
+shakespeare = poet(prompt)
+st.write(shakespeare)
+st.image(poet_image(shakespeare))