import gradio as gr
from diffusers import DiffusionPipeline

def generate_image(steps):
    pipeline = DiffusionPipeline.from_pretrained("nroggendorff/obama", use_safetensors=True)
    pipe = pipeline.to("cuda")
    image = pipe(num_inference_steps=steps).images[0]
    return image

with gr.Blocks() as demo:
    sampling_steps = gr.Slider(value=1000, minimum=20, maximum=1000, label="Sampling Steps", info="How many iterations per image")
    btn = gr.Button("Generate Image")
    output_image = gr.Image(label="Generated Image")
    
    btn.click(fn=generate_image, inputs=sampling_steps, outputs=output_image)

demo.launch()