Create app.py

app.py

@@ -0,0 +1,62 @@
+import streamlit as st
+from diffusers import StableDiffusionPipeline
+from PIL import Image
+import torch
+import io
+import base64
+
+# Model loading (cached)
+@st.cache_resource
+def load_model():
+    pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16).to("cuda" if torch.cuda.is_available() else "cpu")
+    return pipe
+
+pipe = load_model()
+
+st.title("AI Image Generator")
+
+prompt = st.text_area("Enter your prompt:", height=150)
+num_images = st.slider("Number of Images", 1, 4, 1)
+
+if st.button("Generate"):
+    if prompt:
+        with st.spinner("Generating images..."):
+            images = pipe(prompt, num_images=num_images).images
+
+        for i, image in enumerate(images):
+            col1, col2 = st.columns([1, 2])
+
+            with col1:
+                st.image(image, caption=f"Image {i+1}")
+
+            with col2:
+                buffered = io.BytesIO()
+                image.save(buffered, format="PNG")
+                img_str = base64.b64encode(buffered.getvalue()).decode()
+                href = f'<a href="data:image/png;base64,{img_str}" download="image_{i+1}.png">Download Image {i+1}</a>'
+                st.markdown(href, unsafe_allow_html=True)
+
+    else:
+        st.warning("Please enter a prompt.")
+
+# Styling (optional)
+st.markdown("""
+<style>
+body {
+    font-family: sans-serif;
+}
+.stButton>button {
+    background-color: #4CAF50;
+    border: none;
+    color: white;
+    padding: 10px 20px;
+    text-align: center;
+    text-decoration: none;
+    display: inline-block;
+    font-size: 16px;
+    margin: 4px 2px;
+    cursor: pointer;
+    border-radius: 5px;
+}
+</style>
+""", unsafe_allow_html=True)