import gradio as gr
from diffusers import DiffusionPipeline

def generate_image(modelsyu, prompt, negative_prompt):
    pipeline = DiffusionPipeline.from_pretrained(modelsyu)
    pipeline.to("cpu")
    # Attempt to generate an image with the negative prompt if supported
    try:
        image = pipeline(prompt, negative_prompt=negative_prompt).images[0]
    except TypeError:
        # Fallback if negative_prompt is not supported
        image = pipeline(prompt).images[0]
    return image

# Define the Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Text to Image Generation Custom Models Demo")
    prompt = gr.Textbox(label="Prompt", placeholder="Enter your text prompt here")
    negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="Enter your negative prompt here")
    submit_button = gr.Button("Generate Image")
    with gr.Accordion('Load your custom models first'):
        basem = gr.Textbox(label="Your Lora model", value="John6666/pony-diffusion-v6-xl-sdxl-spo")  
    output_image = gr.Image(label="Generated Image")
    submit_button.click(generate_image, inputs=[basem, prompt, negative_prompt], outputs=output_image)

# Launch the demo
demo.launch()