Update app.py

app.py

@@ -3,35 +3,30 @@ from diffusers import DiffusionPipeline
 from PIL import Image
 import torch
 
-# Load the FLUX model
+# Load the diffusion pipeline model
 @st.cache_resource
 def load_pipeline():
-    # Using the 'black-forest-labs/FLUX.1-schnell' model
-    pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16)
-    pipe.enable_model_cpu_offload()  # Offload to CPU to save memory
+    pipe = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
+    pipe.load_lora_weights("Melonie/text_to_image_finetuned")
     return pipe
 
 pipe = load_pipeline()
 
 # Streamlit app
-st.title("Text-to-Image Generation with FLUX.1-schnell")
+st.title("Text-to-Image Generation App")
 
 # User input for prompt
-user_prompt = st.text_input("Enter your image prompt", value="A cat holding a sign that says hello world")
+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..."):
-            # Generate the image using the FLUX model
-            image = pipe(
-                user_prompt,
-                guidance_scale=0.0,  # No guidance
-                num_inference_steps=4,  # Number of steps for faster generation
-            ).images[0]
+            # Generate the image
+            image = pipe(user_prompt).images[0]
             
-            # Save and display the image
-            image.save("generated_image.png")
-            st.image(image, caption="Generated Image (FLUX.1-schnell)", use_column_width=True)
+            # Display the generated image
+            st.image(image, caption="Generated Image", use_column_width=True)
     else:
         st.error("Please enter a valid prompt.")
+