Update app.py

app.py

@@ -1,29 +1,27 @@
 import streamlit as st
-from PIL import Image
-import torch
+from diffusers import DiffusionPipeline
 
-# Optimized lightweight image generation (example)
+# Load the diffusion pipeline model
 @st.cache_resource
-def load_lightweight_model():
-    # Assuming there's a lightweight model for CPU
-    # Placeholder: Replace this with actual lightweight model loading
-    return "lightweight_model"
+def load_pipeline():
+    # Use the 'SaiRaj03/Text_To_Image' model for fast generation
+    pipe = DiffusionPipeline.from_pretrained("SaiRaj03/Text_To_Image")
+    return pipe
 
-model = load_lightweight_model()
+pipe = load_pipeline()
 
 # Streamlit app
-st.title("Text-to-Image Generation (Fast)")
+st.title("Text-to-Image Generation App (Fast)")
 
 # User input for prompt
-user_prompt = st.text_input("Enter your image prompt", value="Astronaut in a jungle")
+user_prompt = st.text_input("Enter your image prompt", value="Astronaut in a jungle, cold color palette, muted colors, detailed, 8k")
 
 # Button to generate the image
 if st.button("Generate Image"):
     if user_prompt:
         with st.spinner("Generating image..."):
-            # Simulated fast image generation (substitute with lightweight model logic)
-            # Example: Generate a simple placeholder image
-            image = Image.new("RGB", (256, 256), color="blue")  # Placeholder image
+            # Generate the image using the new model
+            image = pipe(user_prompt).images[0]
             
             # Display the generated image
             st.image(image, caption="Generated Image", use_column_width=True)