import gradio as gr
from datasets import load_dataset
from diffusers import StableDiffusionPipeline

# Load the dataset
ds = load_dataset("BleachNick/UltraEdit_500k")

# Load the model
pipe = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-refiner-1.0")

def generate_image(prompt, similarity, steps):
    # Adjust the parameters based on the inputs
    pipe.scheduler.set_timesteps(steps)
    guidance_scale = 7.5 * similarity

    # Generate the image
    result = pipe(prompt, guidance_scale=guidance_scale)
    return result.images[0]

# Define the Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Image Generation with Adjustable Parameters")

    with gr.Row():
        prompt = gr.Textbox(label="Prompt")
        similarity = gr.Slider(minimum=0, maximum=1, step=0.01, value=0.5, label="Similarity to Original Image")
        steps = gr.Slider(minimum=1, maximum=100, step=1, value=50, label="Number of Steps")

    with gr.Row():
        generate_button = gr.Button("Generate Image")
        output_image = gr.Image(label="Generated Image")

    generate_button.click(fn=generate_image, inputs=[prompt, similarity, steps], outputs=output_image)

demo.launch()