import os
import shutil
import streamlit as st
from gradio_client import Client

save_dir = r"./textImg/"
if not os.path.exists(save_dir):
    os.makedirs(save_dir)

#Initialize the client
client = Client("black-forest-labs/FLUX.1-schnell")

#Streamlit layout with sidebar for input and main area for image display
st.sidebar.title("Image Generation Prompt")
st.header("🎊Image Generation 🖼️ 🫰")
prompt = st.sidebar.text_area("Enter your prompt:", "")
generate_image = st.sidebar.button("Generate Image")

if generate_image:
    if prompt:
        result = client.predict(
            prompt=prompt,
            seed=0,
            randomize_seed=True,
            width=1024,
            height=1024,
            num_inference_steps=4,
            api_name="/infer"
        )
        
        image_path = result[0]
        save_path = os.path.join(save_dir, "generated_image.jpg")
        shutil.move(image_path, save_path)
        st.image(save_path, caption="Generated Image", use_column_width=True)
        with open(save_path, "rb") as file:
            btn = st.download_button(
                label="Download Image",
                data=file,
                file_name="generated_image.jpg",
                mime="image/jpg"
            )

        st.success(f"Image generated and saved at: {save_path}")
    else:
        st.error("Please enter a prompt to generate the image.")


footer = """
---
#### Made By [Surat Banerjee](https://www.linkedin.com/in/surat-banerjee/) 
"""

st.markdown(footer, unsafe_allow_html=True)