devdeepak/Diffusion_model

1import streamlit as st
2from diffusers import StableDiffusionPipeline
3import torch
4from PIL import Image
5
6# 1. Page Config
7st.set_page_config(page_title="AI Image Generator")
8st.title("🎨 AI Text-to-Image Generator")
9
10# 2. Load the Model
11@st.cache_resource
12def load_model():
13    model_id = "runwayml/stable-diffusion-v1-5"
14    
15    # Check if GPU (CUDA) is available
16    device = "cuda" if torch.cuda.is_available() else "cpu"
17    
18    # Load the pipeline
19    if device == "cuda":
20        # If GPU, use fast float16 precision
21        pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
22    else:
23        # If CPU, use standard float32 precision
24        pipe = StableDiffusionPipeline.from_pretrained(model_id)
25    
26    pipe.to(device)
27    return pipe
28
29st.info("Loading AI Model... The first time may take a few minutes.")
30gen_pipe = load_model()
31
32# 3. User Interface
33prompt = st.text_input("Describe the image you want to see:", placeholder="A futuristic city with flying cars at sunset")
34
35if st.button("Generate Image"):
36    if prompt:
37        with st.spinner("AI is painting... please wait (this takes longer on CPU)"):
38            # Generate the image
39            # num_inference_steps=20 makes it faster for testing; 50 is better quality
40            image = gen_pipe(prompt, num_inference_steps=20).images[0]
41            
42            # Display the image
43            st.image(image, caption=f"Generated: {prompt}", use_column_width=True)
44            
45            # Allow user to download
46            image.save("generated_img.png")
47            with open("generated_img.png", "rb") as file:
48                st.download_button("Download Image", data=file, file_name="ai_image.png", mime="image/png")
49    else:
50        st.warning("Please enter a prompt first!")