import gradio as gr
import torch
from diffusers import DiffusionPipeline

def load_amused_model():
    return DiffusionPipeline.from_pretrained("amused/amused-256")

# Generate image from prompt using AmusedPipeline
def generate_image(prompt):
    try:
        pipe = load_amused_model()
        generator = torch.Generator().manual_seed(8)  # Create a generator for reproducibility
        image = pipe(prompt, generator=generator).images[0]  # Generate image from prompt
        return image, None
    except Exception as e:
        return None, str(e)

def inference(prompt):
    image, error = generate_image(prompt)
    if error:
        return "Error: " + error
    return image

gradio_interface = gr.Interface(
    fn=inference,
    inputs="text",
    outputs="image"
)

if __name__ == "__main__":
    gradio_interface.launch()