Update app.py

app.py

@@ -6,9 +6,14 @@ import torch
 # Load the diffusion pipeline model
 @st.cache_resource
 def load_pipeline():
-    pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0")
-    pipe.load_lora_weights("kothariyashhh/GenAi-Texttoimage")
-    return pipe
+    try:
+        pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16)
+        pipe.load_lora_weights("kothariyashhh/GenAi-Texttoimage")
+        pipe = pipe.to("cuda")  # Move model to GPU if available
+        return pipe
+    except Exception as e:
+        st.error(f"Error loading model: {e}")
+        return None
 
 pipe = load_pipeline()
 
@@ -20,12 +25,18 @@ user_prompt = st.text_input("Enter your image prompt", value="a photo of Yash Ko
 
 # Button to generate the image
 if st.button("Generate Image"):
-    if user_prompt:
+    if user_prompt and pipe:
         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: {e}")
     else:
-        st.error("Please enter a valid prompt.")
+        if not pipe:
+            st.error("Model not loaded. Please check the logs.")
+        else:
+            st.error("Please enter a valid prompt.")