Update app.py

app.py

@@ -1,29 +1,39 @@
 import streamlit as st
-from diffusers import DiffusionPipeline
+from diffusers import FluxPipeline
+from PIL import Image
+import torch
 
-# Load the diffusion pipeline model
+# Load the Flux model
 @st.cache_resource
 def load_pipeline():
-    # Using the 'black-forest-labs/FLUX.1-schnell' model for fast generation
-    pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell")
+    # Using the 'black-forest-labs/FLUX.1-schnell' model
+    pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16)
+    pipe.enable_model_cpu_offload()  # Offload to CPU to save memory
     return pipe
 
 pipe = load_pipeline()
 
 # Streamlit app
-st.title("Text-to-Image Generation App (FLUX Model)")
+st.title("Text-to-Image Generation with FLUX.1-schnell")
 
 # 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="A cat holding a sign that says hello world")
 
 # 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).images[0]
+            image = pipe(
+                user_prompt,
+                guidance_scale=0.0,  # No guidance
+                num_inference_steps=4,  # Number of steps for faster generation
+                max_sequence_length=256,
+                generator=torch.Generator("cpu").manual_seed(0)  # Ensure reproducibility
+            ).images[0]
             
-            # Display the generated image
-            st.image(image, caption="Generated Image (FLUX Model)", use_column_width=True)
+            # Save and display the image
+            image.save("generated_image.png")
+            st.image(image, caption="Generated Image (FLUX.1-schnell)", use_column_width=True)
     else:
         st.error("Please enter a valid prompt.")