from diffusers import DiffusionPipeline, LCMScheduler
import gradio as gr
import torch


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


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

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

iface.launch()