import gradio as gr
import torch
from diffusers import DiffusionPipeline

# Load the BigGAN model using the diffusers library
# The model type is "biggan", which is the correct pipeline to use.
# Move model to GPU if available for faster generation
pipeline = DiffusionPipeline.from_pretrained(
    "osanseviero/BigGAN-deep-128", 
    use_auth_token=True
).to("cuda" if torch.cuda.is_available() else "cpu")

def generate_image(text_input):
    """Generates an image from text using the BigGAN diffusers pipeline."""
    # The pipeline's output is an image.
    image = pipeline(text_input).images[0]
    return image

# Create the Gradio interface directly
interface = gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(label="Enter a prompt"),
    outputs=gr.Image(label="Generated Image"),
    title="BigGAN ImageNet",
    description="BigGAN text-to-image demo.",
    examples=[["american robin"], ["ocean sunset"], ["cat in a hat"]]
)

interface.launch()