# -*- coding: utf-8 -*-
"""Untitled3.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1eg0M7FRTlmONVU_5IaZH-EKADvdtODmT
"""

#!pip install torch

#!pip install diffusers transformers accelerate torch

#!pip install streamlit

from diffusers import StableDiffusionPipeline
import torch
import streamlit as st

def pic_mo(prom):
    model_id = "runwayml/stable-diffusion-v1-5"
    pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
    pipe = pipe.to("cuda")

    prompt = str(prom)

    image = pipe(prompt).images[0]

    return image

def main():
    st.title("Text to Image with runwayml/stable-diffusion-v1-5")

    # Input text prompt
    title = st.text_input('Write a prompt', "")

    if st.button('Generate Image'):
        # Call the pic_mo function to generate an image
        generated_image = pic_mo(title)

        # Display the generated image
        st.image(generated_image, caption='Generated Image', use_column_width=True)

if __name__ == '__main__':
    main()