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

# Loading the model
model_id = "stabilityai/stable-diffusion-xl-base-1.0"
pipe = DiffusionPipeline.from_pretrained(model_id)

# Setting the scheduler
pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)  

# Loading LoRA weights
pipe.load_lora_weights("latent-consistency/lcm-lora-sdxl")

# Converting the model to OpenVINO format
pipeline = OVStableDiffusionXLPipeline.from_pretrained(
    pipe,  
    export=True 
)

def generate_images(prompt, batch_size, num_inference_steps, guidance_scale):
    images = []
    for _ in range(batch_size):
        results = pipeline(
            prompt=prompt, 
            num_inference_steps=num_inference_steps,  
            guidance_scale=guidance_scale  
        )
        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),
        gr.Slider(label="Num Inference Steps", minimum=1, maximum=6, step=1, value=4),
        gr.Slider(label="Guidance Scale", minimum=0.0, maximum=3, step=0.1, value=1.4)
    ],
    outputs=gr.Gallery(label="Generated Images"),
    title="Fast SDXL Generation on CPU"
)

iface.launch()