import os
import gradio as gr
from transformers import AutoModelForCausalLM
from optimum.intel.openvino import OVStableDiffusionPipeline

import torch

# 定義模型 ID 與存儲路徑
model_id = "Kouki79/Realistic_Vision6_LCM"
export_path = "exported_model_openvino_int8"

# 設定圖片大小
HIGH = 1024
WIDTH = 512

print("🔍 檢查 OpenVINO 模型是否已存在...")
if not os.path.exists(export_path) or not os.listdir(export_path):
    print("⚠️ 尚未轉換 OpenVINO 8-bit 模型，開始轉換...")
    
    # 轉換 Hugging Face 模型為 OpenVINO 8-bit
    model = OVStableDiffusionPipeline.from_pretrained(
        model_id,
        export=True,  # 自動轉換為 OpenVINO
        device="CPU",
        precision="int8",  # 啟用 8-bit 量化
    )

    # 儲存轉換後的 OpenVINO 8-bit 模型
    model.save_pretrained(export_path)
    print(f"✅ 轉換完成！OpenVINO 8-bit 模型已儲存至 '{export_path}'")
else:
    print(f"✅ 發現已轉換的 OpenVINO 8-bit 模型：'{export_path}'")

# 載入 OpenVINO 8-bit 模型
print("🔄 載入 OpenVINO 8-bit 模型...")
pipe = OVStableDiffusionPipeline.from_pretrained(
    export_path,
    compile=True,  # 編譯模型以提高效能
    device="CPU",
    safety_checker=None,
    torch_dtype=torch.uint8
)
print("✅ OpenVINO 模型載入完成！")

# 設定推論函數
def infer(prompt):
    print(f"🖼️ 生成圖片: {prompt}")
    image = pipe(
        prompt=f",hyper-realistic 2K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic,",
        negative_prompt="EasyNegative, cartoonish, low resolution, blurry, simplistic, abstract, deformed, ugly,",
        width=WIDTH,
        height=HIGH,
        guidance_scale=1.0,
        num_inference_steps=6,
        num_images_per_prompt=1,
    ).images[0]
    
    return image

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

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(f"""
        # {model_id.split('/')[1]} {WIDTH}x{HIGH}
        Running on OpenVINO (8-bit).
        """)
        
        with gr.Row():
            prompt = gr.Textbox(
                label="Prompt",
                show_label=False,
                max_lines=1,
                placeholder="Enter your prompt",
                container=False,
            )         
            run_button = gr.Button("Generate", scale=0)
        
        result = gr.Image(label="Result", show_label=False)

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

print("🚀 啟動 Gradio Web UI...")
demo.queue().launch()