import gradio as gr
import os
from diffusers import StableDiffusionPipeline
import torch

# Fetch the Hugging Face token from secrets
HF_TOKEN = os.environ.get("HF_TOKEN")

# Load the model
model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(
    model_id,
    use_auth_token=HF_TOKEN
)
pipe = pipe.to("cuda")  # Use GPU

def generate_image(prompt):
    image = pipe(prompt).images[0]
    return image

# Gradio UI
app = gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(label="Describe your image..."),
    outputs=gr.Image(label="Generated Image"),
    title="Free AI Image Generator"
)

app.launch()