import streamlit as st
from diffusers import StableDiffusionPipeline
import torch
import transformers
from transformers import pipeline
pipe = StableDiffusionPipeline.from_pretrained('runwayml/stable-diffusion-v1-5', torch_dtype=torch.float16)


def generate_image(prompt):
   
    image = pipe(prompt).images[0]
    return image

def main():
    st.title("Text-to-Image Generation with Stable Diffusion")

    
    prompt = st.text_input("Enter a prompt for image generation:")

    if st.button("Generate Image"):
        if prompt:
            
            generated_image = generate_image(prompt)

            
            st.image(generated_image, caption="Generated Image", use_column_width=True, channels="RGB")

            
            generated_image.save("generated_image.png")

           
            st.markdown(get_image_download_link(generated_image), unsafe_allow_html=True)

def get_image_download_link(image):
    """Generate a link allowing the image to be downloaded"""
    href = f'<a href="generated_image.png" download>Click here to download the image</a>'
    return href

if __name__ == "__main__":
    main()