Create app.py

app.py

@@ -0,0 +1,192 @@
+import streamlit as st
+import torch
+from diffusers import StableDiffusionPipeline
+from PIL import Image
+import numpy as np
+
+# Set page config
+st.set_page_config(
+    page_title="AI Image Generator",
+    page_icon="🎨",
+    layout="centered"
+)
+
+# Cache the model loading to avoid reloading on every interaction
+@st.cache_resource
+def load_model():
+    """Load and cache the Stable Diffusion model"""
+    model_id = "runwayml/stable-diffusion-v1-5"
+    
+    # Check if CUDA is available
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    
+    # Load the pipeline
+    pipe = StableDiffusionPipeline.from_pretrained(
+        model_id,
+        torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+        use_safetensors=True
+    )
+    pipe = pipe.to(device)
+    
+    # Enable memory efficient attention if using CUDA
+    if device == "cuda":
+        pipe.enable_attention_slicing()
+        pipe.enable_memory_efficient_attention()
+    
+    return pipe
+
+def generate_image(prompt, negative_prompt="", num_inference_steps=20, guidance_scale=7.5, width=512, height=512):
+    """Generate image from text prompt"""
+    try:
+        pipe = load_model()
+        
+        # Generate image
+        with torch.no_grad():
+            image = pipe(
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                num_inference_steps=num_inference_steps,
+                guidance_scale=guidance_scale,
+                width=width,
+                height=height
+            ).images[0]
+        
+        return image
+    except Exception as e:
+        st.error(f"Error generating image: {str(e)}")
+        return None
+
+def main():
+    # Header
+    st.title("🎨 AI Image Generator")
+    st.markdown("Generate beautiful images from text descriptions using Stable Diffusion!")
+    
+    # Sidebar for advanced settings
+    with st.sidebar:
+        st.header("⚙️ Settings")
+        
+        # Image dimensions
+        col1, col2 = st.columns(2)
+        with col1:
+            width = st.selectbox("Width", [512, 768, 1024], index=0)
+        with col2:
+            height = st.selectbox("Height", [512, 768, 1024], index=0)
+        
+        # Generation parameters
+        num_inference_steps = st.slider("Inference Steps", 10, 50, 20, 
+                                       help="More steps = better quality but slower")
+        guidance_scale = st.slider("Guidance Scale", 1.0, 20.0, 7.5, 0.5,
+                                  help="Higher values = more adherence to prompt")
+        
+        # Info
+        st.markdown("---")
+        st.markdown("### 💡 Tips")
+        st.markdown("- Be specific in your descriptions")
+        st.markdown("- Use artistic styles (e.g., 'oil painting', 'digital art')")
+        st.markdown("- Add quality modifiers (e.g., 'highly detailed', '4k')")
+        st.markdown("- Use negative prompts to avoid unwanted elements")
+    
+    # Main content area
+    col1, col2 = st.columns([2, 1])
+    
+    with col1:
+        # Text input for prompt
+        prompt = st.text_area(
+            "✍️ Describe the image you want to generate:",
+            placeholder="A beautiful sunset over mountains, oil painting style, highly detailed",
+            height=100
+        )
+        
+        # Negative prompt (optional)
+        negative_prompt = st.text_area(
+            "❌ Negative prompt (optional - things to avoid):",
+            placeholder="blurry, low quality, distorted",
+            height=60
+        )
+        
+        # Generate button
+        generate_btn = st.button("🚀 Generate Image", type="primary", use_container_width=True)
+    
+    with col2:
+        # Example prompts
+        st.markdown("### 🎯 Example Prompts")
+        examples = [
+            "A majestic lion in a savanna at sunset",
+            "Cyberpunk cityscape at night, neon lights",
+            "Van Gogh style painting of a coffee shop",
+            "Cute robot playing with cats in a garden",
+            "Abstract art with vibrant colors and geometric shapes"
+        ]
+        
+        for i, example in enumerate(examples):
+            if st.button(f"Use Example {i+1}", key=f"example_{i}"):
+                st.session_state.example_prompt = example
+        
+        # Apply example if selected
+        if hasattr(st.session_state, 'example_prompt'):
+            prompt = st.session_state.example_prompt
+            del st.session_state.example_prompt
+            st.rerun()
+    
+    # Generate and display image
+    if generate_btn and prompt:
+        with st.spinner("🎨 Creating your masterpiece... This may take a few moments!"):
+            # Show progress
+            progress_bar = st.progress(0)
+            for i in range(100):
+                progress_bar.progress(i + 1)
+                if i == 99:
+                    break
+            
+            # Generate image
+            image = generate_image(
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                num_inference_steps=num_inference_steps,
+                guidance_scale=guidance_scale,
+                width=width,
+                height=height
+            )
+            
+            progress_bar.empty()
+        
+        if image:
+            # Display the generated image
+            st.success("✅ Image generated successfully!")
+            st.image(image, caption=f"Generated from: '{prompt}'", use_column_width=True)
+            
+            # Download button
+            img_buffer = io.BytesIO()
+            image.save(img_buffer, format='PNG')
+            st.download_button(
+                label="📥 Download Image",
+                data=img_buffer.getvalue(),
+                file_name=f"generated_image_{hash(prompt) % 10000}.png",
+                mime="image/png",
+                use_container_width=True
+            )
+            
+            # Show generation parameters
+            with st.expander("📊 Generation Details"):
+                st.json({
+                    "prompt": prompt,
+                    "negative_prompt": negative_prompt,
+                    "dimensions": f"{width}x{height}",
+                    "inference_steps": num_inference_steps,
+                    "guidance_scale": guidance_scale
+                })
+    
+    elif generate_btn and not prompt:
+        st.warning("⚠️ Please enter a prompt to generate an image!")
+    
+    # Footer
+    st.markdown("---")
+    st.markdown(
+        "Built with ❤️ using [Streamlit](https://streamlit.io) and "
+        "[Stable Diffusion](https://huggingface.co/runwayml/stable-diffusion-v1-5)"
+    )
+
+if __name__ == "__main__":
+    # Add missing import
+    import io
+    main()