import gradio as gr
from diffusers import DiffusionPipeline

# Load the pre-trained model
pipe = DiffusionPipeline.from_pretrained("John6666/mala-anime-mix-nsfw-pony-xl-v3-sdxl")

# Function to generate image based on the input prompt
def generate_image(prompt):
    image = pipe(prompt).images[0]
    return image

# Create the Gradio interface using Blocks
with gr.Blocks() as demo:
    gr.Markdown("# Image Generation with DiffusionPipeline")
    
    with gr.Row():
        prompt_input = gr.Textbox(label="Prompt", placeholder="Enter your prompt here")
    
    with gr.Row():
        generate_btn = gr.Button("Generate Image")
    
    with gr.Row():
        output_image = gr.Image(label="Generated Image")
    
    # Define the button click event
    generate_btn.click(fn=generate_image, inputs=prompt_input, outputs=output_image)

# Launch the Gradio app
demo.launch()