Text_To_Image / app.py
shubham5027's picture
Update app.py
5d68ba9
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()