from diffusers import DiffusionPipeline, LCMScheduler
from optimum.intel import OVStableDiffusionXLPipeline
import gradio as gr


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

pipeline = OVStableDiffusionXLPipeline.from_pretrained(
    pipe,
    export=True 
)

def generate_images(prompt, batch_size):
    images = []
    for _ in range(batch_size):
        results = pipeline(prompt)
        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 Generation on CPU"
)

iface.launch(debug=True)