Update app.py

app.py

@@ -1,5 +1,6 @@
 import streamlit as st
 from diffusers import DiffusionPipeline
+from peft import PeftModel  # Import PEFT model loader
 from PIL import Image
 import torch
 
@@ -7,7 +8,16 @@ import torch
 @st.cache_resource
 def load_pipeline():
     pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
-    pipe.load_lora_weights("Melonie/text_to_image_finetuned")
+    
+    # Ensure PEFT is available before loading LoRA weights
+    try:
+        pipe.load_lora_weights("Melonie/text_to_image_finetuned")
+    except ValueError as e:
+        st.error("PEFT backend is required but not properly set up.")
+        raise e
+    
+    if torch.cuda.is_available():
+        pipe.to("cuda")  # Move pipeline to GPU if available
     return pipe
 
 pipe = load_pipeline()
@@ -16,17 +26,23 @@ pipe = load_pipeline()
 st.title("Text-to-Image Generation App")
 
 # User input for prompt
-user_prompt = st.text_input("Enter your image prompt", value="Astronaut in a jungle, cold color palette, muted colors, detailed, 8k")
+user_prompt = st.text_input(
+    "Enter your image prompt", 
+    value="Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
+    help="Provide a detailed description of the image you'd like to generate."
+)
 
 # Button to generate the image
 if st.button("Generate Image"):
     if user_prompt:
         with st.spinner("Generating image..."):
-            # Generate the image
-            image = pipe(user_prompt).images[0]
-            
-            # Display the generated image
-            st.image(image, caption="Generated Image", use_column_width=True)
+            try:
+                # Generate the image
+                image = pipe(user_prompt).images[0]
+                
+                # Display the generated image
+                st.image(image, caption="Generated Image", use_column_width=True)
+            except Exception as e:
+                st.error(f"Error generating image: {str(e)}")
     else:
         st.error("Please enter a valid prompt.")
-