# Streamlit and Machine Learning libraries
import streamlit as st
import torch
from torch import autocast
from diffusers import StableDiffusionPipeline

# Libraries for processing image
from PIL import Image

# Private modules
from authtoken import auth_token


# Download stable diffusion model from Hugging Face
modelid = "CompVis/stable-diffusion-v1-4"
stable_diffusion_model = StableDiffusionPipeline.from_pretrained(modelid, revision="fp16", dtype=torch.float16, use_auth_token=auth_token)

# Create a Streamlit app
st.set_page_config(
    page_title="Text to Image App",
    page_icon="🖼️",
    layout="centered",
)

# Create input box on the user interface
st.write("# Text to Image app")
prompt = st.text_area("Enter your text here:", height=10, max_chars=200)

# Create a placeholder to show the generated image
img_placeholder = st.empty()

# Generate image from text
def generate_image():
    if prompt:
        st.write("Generating image...")
        try:
            with autocast():
                image = stable_diffusion_model(prompt, guidance_scale=8.5)["sample"][0]

            # Display the generated image on the user interface
            st.image(image, caption="Generated Image", use_column_width=True)
        except Exception as e:
            st.error(f"Error generating the image: {str(e)}")

# Create a button to trigger image generation
if st.button("Generate Image"):
    generate_image()