import gradio as gr
import torch
from diffusers import StableDiffusion3Pipeline
import os
from huggingface_hub import InferenceApi


# Retrieve Hugging Face API key from environment variables
hf_api_key = os.getenv("PRODIGY_GA_02")

# Load model with authentication token
model_id = "stabilityai/stable-diffusion-3.5-medium"
pipe = StableDiffusion3Pipeline.from_pretrained("stabilityai/stable-diffusion-3.5-medium")

pipe.to("cuda")# If you have GPU access; otherwise, use "cpu"

# Define Gradio interface function
def generate_image(prompt):
    images = pipe(prompt).images
    return images[0]

# Create Gradio UI
iface = gr.Interface(
    fn=generate_image,
    inputs="text",
    outputs="image",
    title="Stable Diffusion Generator",
    description="Enter a text prompt to generate an image",
)

# Launch the interface
iface.launch()