import gradio as gr
import openvino.runtime as ov
from optimum.intel.openvino import OVStableDiffusionPipeline

model_id = "HelloSun/chilloutmix_NiPrunedFp32Fix-openvino"

# 確保這些是有效的尺寸
HIGH = 2048
WIDTH = 1024

pipe = OVStableDiffusionPipeline.from_pretrained(model_id)

negative_prompt = "(worst quality, low quality, lowres), zombie, interlocked fingers,"

def infer(prompt, negative_prompt):
    image = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        width=WIDTH,  # 使用 WIDTH
        height=HIGH,  # 使用 HIGH
        guidance_scale=7.5,
        num_inference_steps=28,
        num_images_per_prompt=1,
    ).images[0]
    return image

examples = [
    "1girl,hakurei reimu,bird,flower,hair tubes,solo,bow,hair bow,bird on head,brown hair,long hair,looking at viewer,animal on head,branch,on head,red bow,purple eyes,lips,",
    "1girl, silver hair, symbol-shaped pupils, yellow eyes, smiling, light particles, light rays, wallpaper, star guardian, serious face, red inner hair, power aura, grandmaster1, golden and white clothes",
    "A cute kitten, Tinkle style.",
    "(illustration, 8k CG, extremely detailed),(whimsical),catgirl,teenage girl,playing in the snow,winter wonderland,snow-covered trees,soft pastel colors,gentle lighting,sparkling snow,joyful,magical atmosphere,highly detailed,fluffy cat ears and tail,intricate winter clothing,shallow depth of field,watercolor techniques,close-up shot,slightly tilted angle,fairy tale architecture,nostalgic,playful,winter magic,(masterpiece:2),best quality,ultra highres,original,extremely detailed,perfect lighting,",
]

css = """
#col-container {
    margin: 0 auto;
    max-width: 520px;
}
"""

power_device = "CPU"

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(f"""
        # chilloutmix_NiPrunedFp32Fix-openvino  {HIGH}x{WIDTH}
        Currently running on {power_device}.
        """)
        
        with gr.Row():
            prompt_input = gr.Text(
                label="Prompt",
                show_label=False,
                max_lines=1,
                placeholder="Enter your prompt",
                container=False,
            )         
            run_button = gr.Button("Run", scale=0)
        
        result = gr.Image(label="Result", show_label=False)

        gr.Examples(
            examples=examples,
            fn=infer,
            inputs=[prompt_input],
            outputs=[result]
        )

    run_button.click(
        fn=infer,
        inputs=[prompt_input],
        outputs=[result]
    )

demo.queue().launch()