Create app.py

app.py

@@ -0,0 +1,56 @@
+import streamlit as st
+import torch
+from diffusers import FluxPipeline
+import io
+
+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")
+
+# Load the model with caching to avoid reloading on every interaction
+@st.cache_resource
+def load_pipeline():
+    pipe = FluxPipeline.from_pretrained(
+        "black-forest-labs/FLUX.1-schnell", 
+        torch_dtype=torch.bfloat16
+    )
+    # Use CPU offload to save memory on Hugging Face's free tier
+    pipe.enable_model_cpu_offload()
+    return pipe
+
+pipeline = load_pipeline()
+
+# Sidebar for 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)
+
+# User Input
+prompt = st.text_area("Enter your prompt:", "A futuristic mechanical engine, detailed blueprint style, blue neon lines")
+
+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]
+            
+            # 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")
+    else:
+        st.warning("Please enter a prompt first!")