import os
import torch
import gradio as gr
from diffusers import OnnxStableDiffusionPipeline

MODEL_ID = "SimianLuo/LCM_Dreamshaper_v7"
ONNX_PATH = os.getenv("ONNX_PATH", "./onnx-model")

pipe = OnnxStableDiffusionPipeline.from_pretrained(
    ONNX_PATH,
    provider="CUDAExecutionProvider" if torch.cuda.is_available() else "CPUExecutionProvider",
)

def generate(prompt, negative_prompt, steps, guidance_scale, width, height, seed):
    generator = torch.Generator(device="cuda" if torch.cuda.is_available() else "cpu")
    if seed >= 0:
        generator = generator.manual_seed(seed)

    image = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt or None,
        num_inference_steps=steps,
        guidance_scale=guidance_scale,
        width=width,
        height=height,
        generator=generator,
    ).images[0]

    return image

demo = gr.Blocks()
with demo:
    gr.Markdown("# LCM DreamShaper v7 ONNX")
    with gr.Row():
        prompt = gr.Textbox(label="Prompt", value="a cinematic portrait of a futuristic astronaut")
    negative_prompt = gr.Textbox(label="Negative prompt", value="blurry, low quality")
    with gr.Row():
        steps = gr.Slider(1, 8, value=4, step=1, label="Steps")
        guidance_scale = gr.Slider(0, 10, value=1.5, step=0.1, label="Guidance scale")
    with gr.Row():
        width = gr.Slider(256, 1024, value=512, step=64, label="Width")
        height = gr.Slider(256, 1024, value=512, step=64, label="Height")
        seed = gr.Number(value=42, precision=0, label="Seed")
    btn = gr.Button("Generate")
    out = gr.Image(label="Result")

    btn.click(
        fn=generate,
        inputs=[prompt, negative_prompt, steps, guidance_scale, width, height, seed],
        outputs=out,
    )

demo.launch()