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

# Load Hugging Face access token from secrets
hf_token = os.getenv("secret")  # Ensure your secret is named "secret"

# Set up the pipeline with the access token
pipeline = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    use_auth_token=hf_token
).to("cuda" if torch.cuda.is_available() else "cpu")

# Inference function
def generate_image(prompt):
    with torch.no_grad():
        image = pipeline(prompt).images[0]
    return image

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# FLUX Image Generator")
    prompt = gr.Textbox(label="Enter your prompt", placeholder="e.g. Astronaut riding a horse")
    generate_btn = gr.Button("Generate Image")
    output_image = gr.Image(label="Generated Image")

    generate_btn.click(fn=generate_image, inputs=prompt, outputs=output_image)

# Launch the app
demo.launch()