from diffusers import DiffusionPipeline, LCMScheduler
import gradio as gr


pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0").to("cuda") 
pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
pipe.load_lora_weights("latent-consistency/lcm-lora-sdxl")


def generate_images(prompt, batch_size):
    images = []
    for _ in range(batch_size):
        results = pipe(
            prompt=prompt,
            num_inference_steps=4,
            guidance_scale=1.4,
        )
        images.append(results.images[0])
    return images

iface = gr.Interface(
    fn=generate_images,
    inputs=[
        gr.Textbox(label="Prompt"),
        gr.Slider(label="Batch Size", minimum=1, maximum=12, step=1, value=1)
    ],
    outputs=gr.Gallery(label="Generated Images"),
    title="SuperFast SDXL Generation."
)

iface.launch()