import os
import streamlit as st
from diffusers import StableDiffusionPipeline
import torch
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Hugging Face API key from environment
HUGGING_FACE_API_KEY = os.getenv("HUGGING_FACE_API_KEY")

# Load Hugging Face model using API key
model_id = "stabilityai/stable-diffusion-xl-base-1.0"
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = StableDiffusionPipeline.from_pretrained(model_id, use_auth_token=HUGGING_FACE_API_KEY).to(device)

# Streamlit UI setup
st.title("Text-to-Image Generator using Hugging Face and Stable Diffusion")
st.write("Generate images based on text descriptions using the Hugging Face model.")

# Text input
text_prompt = st.text_input("Enter your text prompt:", "")

# Generate image on button click
if st.button("Generate Image"):
    if text_prompt:
        with st.spinner("Generating image..."):
            try:
                image = pipe(text_prompt).images[0]
                image_path = os.path.join("output", "generated_image.png")
                image.save(image_path)
                st.image(image, caption="Generated Image")
            except Exception as e:
                st.error(f"Error generating image: {str(e)}")
    else:
        st.warning("Please enter a text prompt.")

# Ensure output directory exists
if not os.path.exists('output'):
    os.makedirs('output')