Update app.py

app.py

@@ -2,55 +2,73 @@ import streamlit as st
 import torch
 from diffusers import FluxPipeline
 import io
+import os
 
+# Page configuration
 st.set_page_config(page_title="Flux Image Generator", layout="centered")
 
 st.title("🎨 AI Image Generator")
-st.caption("Powered by FLUX.1 [schnell] - Optimized for speed and quality")
+st.caption("Powered by FLUX.1 [schnell]")
 
-# Load the model with caching to avoid reloading on every interaction
+# 1. THE LOAD FUNCTION
 @st.cache_resource
 def load_pipeline():
+    # This pulls the secret you named 'HF_TOKEN' from your Space Settings
+    token = os.getenv("HF_TOKEN")
+    
+    # Loading the model with the access token
     pipe = FluxPipeline.from_pretrained(
         "black-forest-labs/FLUX.1-schnell", 
-        torch_dtype=torch.bfloat16
+        torch_dtype=torch.bfloat16,
+        token=token
     )
-    # Use CPU offload to save memory on Hugging Face's free tier
+    
+    # Enables memory saving for the Hugging Face free tier
     pipe.enable_model_cpu_offload()
     return pipe
 
-pipeline = load_pipeline()
+# Initialize the pipeline
+try:
+    pipeline = load_pipeline()
+except Exception as e:
+    st.error("Could not load the model. Make sure you accepted the terms on the model page and added your HF_TOKEN to secrets.")
+    st.stop()
 
-# Sidebar for settings
+# 2. SIDEBAR SETTINGS
 with st.sidebar:
     st.header("Settings")
     width = st.slider("Width", 512, 1024, 1024, step=128)
     height = st.slider("Height", 512, 1024, 1024, step=128)
     num_steps = st.slider("Inference Steps", 1, 4, 4)
+    st.info("Tip: Use 4 steps for the best quality.")
 
-# User Input
-prompt = st.text_area("Enter your prompt:", "A futuristic mechanical engine, detailed blueprint style, blue neon lines")
+# 3. USER INTERFACE
+prompt = st.text_area("Enter your prompt:", "A futuristic robotic arm building a circuit board, cinematic lighting, 8k")
 
 if st.button("Generate Image"):
     if prompt:
-        with st.spinner("Generating... please wait."):
-            # Generate the image
-            image = pipeline(
-                prompt,
-                num_inference_steps=num_steps,
-                guidance_scale=0.0,
-                width=width,
-                height=height,
-                max_sequence_length=256
-            ).images[0]
+        with st.spinner("Generating... this may take a minute."):
+            try:
+                # Generate the image
+                image = pipeline(
+                    prompt,
+                    num_inference_steps=num_steps,
+                    guidance_scale=0.0,
+                    width=width,
+                    height=height,
+                    max_sequence_length=256
+                ).images[0]
+                
+                # Display image
+                st.image(image, caption="Generated Result", use_column_width=True)
+                
+                # Download button
+                buf = io.BytesIO()
+                image.save(buf, format="PNG")
+                byte_im = buf.getvalue()
+                st.download_button(label="Download Image", data=byte_im, file_name="generated.png", mime="image/png")
             
-            # Display image
-            st.image(image, caption="Generated Result", use_column_width=True)
-            
-            # Download button
-            buf = io.BytesIO()
-            image.save(buf, format="PNG")
-            byte_im = buf.getvalue()
-            st.download_button(label="Download Image", data=byte_im, file_name="generated.png", mime="image/png")
+            except Exception as e:
+                st.error(f"Error during generation: {e}")
     else:
         st.warning("Please enter a prompt first!")