# -*- coding: utf-8 -*-

#!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.float32)  # Change to float32
    # No need ("cuda") if running on CPU because dont't have money 

    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 (จะใช้เวลาค่อนข้างมากในการสร้างภาพเนื่องจากใช้ CPU ในการรันโมเดล)', "")

    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()