File size: 1,136 Bytes
5c4a3bd
 
 
 
 
 
5d68ba9
5c4a3bd
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
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()